Custom loss with instance weights

G’day and apologies for any mistakes (first post here…)!

I’ve set up an object detector CNN that (for now) predicts the label of cells on a WxH grid (background vs. object).
So far I’ve used pyTorch’s integrated Cross-Entropy loss, but unfortunately this has no implementation of instance weights. My task heavily relies on individual weighting (for example, I need to disable backprop. in the vicinity of an object to allow the model to have some translation shifts), and therefore I need a solution for that.

Ideally I’d like to have an instance-weighted multi-task loss (cross-entropy for the class, regression for bounding box coordinates), but to start simple let’s ignore the coordinates. I tried to create a drop-in layer before the loss that just multiplies the gradients with a weight grid during the backward pass. Something like this:

class WeightMultiplier(Function):

   def forward(self, input, target, instanceWeights):
       # Technically I'd need to apply the weights also during the forward pass,
       # but for simplicity reasons I didn't intend to re-implement a weighted
       # cross-entropy loss.
       # So I'll just do an identity transform here.
       
       self.instanceWeights = instanceWeights
       
       return input
       
       
       
   def backward(self, grad_output):
       # Here I'd like to have the grad_output multiplied by the instance weights,
       # but this step is never called...
       
       return torch.mul(grad_output,self.instanceWeights)

However, if I insert this layer in my model:
pred_label = model(imgTensor)
pred_label = model.weights_grid(pred_label,labelTensor)
loss_label = model.label_loss(pred_label,labelTensor.squeeze())
loss_label.backward()
optimizer.step()

the backward pass is never executed.
I’m positive there’s lots of the autograd functionality I still have to learn. So what am I doing wrong? And how would you proceed with something like this (instance-weighted multi-task loss)?

Many thanks for your answers.

– EDIT –
I don’t get it. I’ve updated some of the software packages but didn’t change anything else. Now suddenly only one of my GPUs work (the other freezes PyDev), but out of the blue the backward pass actually gets called!
So in this case I guess I can re-formulate my question: what is the best practice to implement a multi-task instance-weighted loss in PyTorch?

Correct me if I’m wrong, but as I understand it, you’re looking for a way to have a loss function (something like CrossEntropy loss) weigh each instance separately. I’m currently working on something that will have loss functions return a loss per instance, so you can apply weights to these losses and sum them together yourself.

If you want to implement an instance-weighted loss, I think you’d have to re-implement a loss.