Extracting weights from SGD algorithm

Question

So I am implementing SGD for a binary classification problem. There are 2 classes of points and I want to plot the decision boundary but I'm not sure how to extract the weights from the code to plot it.

Here is the code:

def train_model(train_dl, model):
    # define the optimization
    criterion = nn.BCELoss(reduction='none')
    optimizer = torch.optim.SGD(net.parameters(), lr=0.1)
    # enumerate epochs
    for epoch in range(10):
        # enumerate mini batches
        for i, (inputs, targets) in enumerate(train_dl):
            # clear the gradients
            optimizer.zero_grad()
            # compute the model output
            yhat = model(inputs)
            # calculate loss
            loss = criterion(yhat, targets)
            # credit assignment
            loss.backward()
            # update model weights
            optimizer.step()

Any help would be much appreciated! Thanks.

Answer 1

You do not extract this information from the SGD optimizer, this information is part of your model.
What you can do, at test time, is generate a grid of points, compute their prediction using the trained model and then plot the grid points coloring them according to the prediction.