nn.MSELoss inputs inquiry

1

Hello,
I have a question with regards to nn.MSELoss(). I’m trying to implement a network to demosaic/debayer an image. I have seen that the inputs to nn.MSELoss should be nn.MSELoss(outputs, target). My outputs tensor is of shape [batch_size, 3, 224, 224], but my target is [batch_size].

I am loading the images as so.

trainset = datasets.ImageFolder(root=trainpath,transform=preprocess)
trainloader = DataLoader(dataset=trainset,batch_size=batch_size, shuffle=True)

And then in my training function

for data in trainloader:
    inputs, target = data[0], data[1]
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = nn.MSELoss(outputs, target)

Should it instead be loss = nn.MSELoss(outputs, inputs) as I know MSE subtracts the two values. But from everything I’ve seen, it’s always (outputs, target).

Thank you very much

2

As you’ve explained, nn.MSELoss expects the model output and target to have the same shape.
Since your target shape differs, you would have to explain a bit more how the loss calculation should be performed for your use case.
E.g. you could repeat the target values for each channel and pixel, but I doubt this would fit your use case.