Custom loss_func and get_preds with_loss=True

I have been playing around with custom loss functions and then wanted to use fastai’s get_preds(with_loss=True) but found it doesn’t work with my implementation. Is there something special I need to add/do to my loss function to have it work to give me the loss for each sample? I noticed that the built in losses like F.mse_loss work but actually give the loss per output node as opposed to an aggregated calculation across all of them. Mine aggregates all outputs so I expect this is the reason, however it works to train the model. Is it possible to output the aggregated version? Or change mine to work?

My loss function:

class mse_custom(nn.Module):
    def __init__(self,weight,order):
        super().__init__()
        self.weight = weight
        self.order = torch.Tensor(order)
        
    def forward(self, preds, target):
        assert not target.requires_grad
        assert preds.size(0) == target.size(0)
        abs_losses = torch.abs(preds - target)**2
        abs_weighted_losses = abs_losses*(torch.abs(target)*self.weight)*self.order.cuda()
        loss = torch.mean(abs_weighted_losses)
        return loss

The error msg:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-1283-2b985cd61498> in <module>
----> 1 valid_pred,valid_targ,valid_loss = learn.get_preds(ds_type=DatasetType.Valid,with_loss=True)
      2 valid_preds = (valid_pred,valid_targ)

~/anaconda3/lib/python3.6/site-packages/fastai/basic_train.py in get_preds(self, ds_type, with_loss, n_batch, pbar)
    332         lf = self.loss_func if with_loss else None
    333         return get_preds(self.model, self.dl(ds_type), cb_handler=CallbackHandler(self.callbacks),
--> 334                          activ=_loss_func2activ(self.loss_func), loss_func=lf, n_batch=n_batch, pbar=pbar)
    335 
    336     def pred_batch(self, ds_type:DatasetType=DatasetType.Valid, batch:Tuple=None, reconstruct:bool=False) -> List[Tensor]:

~/anaconda3/lib/python3.6/site-packages/fastai/basic_train.py in get_preds(model, dl, pbar, cb_handler, activ, loss_func, n_batch)
     43            zip(*validate(model, dl, cb_handler=cb_handler, pbar=pbar, average=False, n_batch=n_batch))]
     44     if loss_func is not None:
---> 45         with NoneReduceOnCPU(loss_func) as lf: res.append(lf(res[0], res[1]))
     46     if activ is not None: res[0] = activ(res[0])
     47     return res

~/anaconda3/lib/python3.6/site-packages/fastai/torch_core.py in __enter__(self)
    278     def __enter__(self):
    279         if hasattr(self.loss_func, 'weight') and self.loss_func.weight is not None:
--> 280             self.device = self.loss_func.weight.device
    281             self.loss_func.weight = self.loss_func.weight.cpu()
    282         if hasattr(self.loss_func, 'reduction'):

AttributeError: 'int' object has no attribute 'device'

Thanks for the help

Ah, figured it out. In case anyone else runs into this there were two issues with my code.

  1. I overwrote the weight attribute used by other loss functions… so i changed it to w to fix that issue.
  2. As I suspected the reduction has to be baked into the loss function. It also turns out that the losses are computed on the cpu so you need to make the computation occur with non-gpu tensors.

The new loss function that works:

class mse_custom(nn.Module):
    def __init__(self,weight,order):
        super().__init__()
        self.w = weight
        self.order = torch.Tensor(order)
        
    def forward(self, preds, target, reduction='mean'):
        assert not target.requires_grad
        assert preds.size(0) == target.size(0)
        abs_losses = torch.abs(preds - target)**2
        if reduction == 'mean':
            return torch.mean(abs_losses*(torch.abs(target)*self.w)*self.order.cuda())
        else:
            return abs_losses*(torch.abs(target)*self.w)*self.order