Loss function on output slices

hello

I’m doing regression and my network output is a 14 values vector

by using MSE loss function the training converge but I’d like to use different criteria for the different values

as a starter, I tried something like this

lass MyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss()

    def forward(self, yhat, y):
        loss = 2*self.mse(yhat[:2], y[:2]) + 5*self.mse(yhat[2:5], y[2:5]) + self.mse(yhat[5:], y[5:])
        return loss

however this results in (I set batch size to 2 for testing purposes)

UserWarning: Using a target size (torch.Size([0, 14])) that is different to the input size (torch.Size([2, 14])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.

I reckon I’m not slicing the tensor correctly, or is this completely the wrong way to go about this?

The alternative I have in mind is to have multiple output layers with a single loss function each; is that a better idea?

class MyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss()

    def forward(self, yhat, y):
        loss = 2*self.mse(yhat[:, :2], y[:, :2]) + 5*self.mse(yhat[:, 2:5], y[:, 2:5]) + self.mse(yhat[:, 5:], y[:, 5:])
        return loss
1 Like

I wonder if slicing the input of the loss function is ok to do that? This will not affect the autograd in PyTorch? In fact, I’m having the similar problem that I need to slice the input and feed it to a customized loss function.

Y = Model(input) # Y is a batch_size by n by n matrix
loss = loss_1(Y[:, :10, :10]) + loss_2(Y[:, 10:, 10:])

Thanks