How to ignore loss only some of the time in PyTorch

user_1828289 0
Viewed 476

While training a neural net in PyTorch, I'd like to be able to look at the model outputs and decide whether a particular output should result in a loss or just be ignored (=zero loss).

The usual training loop looks something like this:

for (data, labels) in loader:
    data, labels = data.to(device), labels.to(device)
    optimizer.zero_grad()
    logits = model(data)
    loss = criterion(logits, labels)
    loss.backward()
    optimizer.step()

What is the best way to zero out some of the losses?

Should I make some of the labels equal to the logits before calculating criterion?

Or, can I modify the loss after calculating it before calling loss.backward()? (and how? I could multiply it by a mask of zeros and ones, or round-trip it to a numpy array and modify it, etc)

The use case here is semantic segmentation, sort of. The labels are a coarse image, where each region (corresponding to a 32x32 patch in the input image) can be either "yes", "no", or "ambiguous". The "ambiguous" ones should not produce a loss, whatever the network output is. In this case I know which losses should be ignored right up front, from the labels; but I'd like to also know how to ignore based on the outputs as well as the labels.

conv-neural-network pytorch neural-network
1 Answers

While training the network, you can save the best model if validation loss as decreased. Try this code:

if valid_loss <= valid_loss_min:
        print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(valid_loss_min, valid_loss))
        torch.save(model.state_dict(), 'model_name.pt')
        valid_loss_min = valid_loss
user_132122690
Related