Built in way to only update weights corresponding to certain outputs

2

Hi,

The proposed method is actually the right way to do it.
Another approach not using hooks is to no include these rows into the computation of the loss, so that the .backward() will not required to be hooked into.
Few examples below:

output = model(data)
to_keep_indices = get_indices_to_keep(output)
output = output.index_select(0, to_keep_indices)
target = target.index_select(0, to_keep_indices)
# You can also do the same with a 0/1 mask and the masked_select function
loss = criterion(output, target)

Or

# Create your criterion as:
criterion = nn.NLLLoss(reduce=False) # EDITED: wrong keyword argument

# In your training loop
loss = criterion(output, target)
to_keep_indices = get_indices_to_keep(output, loss)
final_loss = loss.index_select(0, to_keep_indices)
# You can also do the same with a 0/1 mask and the masked_select function
# Now average the loss of each of the samples you want to keep and backward.
final_loss.mean().backward()
1 Like