Pytorch: Dimensions for cross entropy is correct but somehow wrong for MSE?

Viewed 820

I was creating a program that would take in as input the Fashion MNIST set and I was tweaking around with my model to see how different parameters would change the accuracy.

One of the tweaks I made to my model was to change my model's loss function from cross entropy to MSE.

# The code above is miscellaneous training data import code

trainloader = torch.utils.data.DataLoader(trainset, batch_size = 64, shuffle = True, num_workers=4)
testloader = torch.utils.data.DataLoader(testset, batch_size = 64, shuffle = True, num_workers=4)

dataiter = iter(trainloader)
images, labels = dataiter.next()
from torch import nn, optim
import torch.nn.functional as F

model = nn.Sequential(nn.Linear(784, 128),
                      nn.ReLU(),
                      nn.Linear(128, 10),
                      nn.LogSoftmax(dim = 1)
                     )
model.to(device)

# Define the loss
criterion = nn.MSELoss()

# Define the optimizer
optimizer = optim.Adam(model.parameters(), lr = 0.001)

# Define the epochs
epochs = 5

train_losses, test_losses = [], []

for e in range(epochs):
  running_loss = 0
  for images, labels in trainloader:
    # Flatten Fashion-MNIST images into a 784 long vector
    images = images.to(device)
    labels = labels.to(device)
    images = images.view(images.shape[0], -1)

    # Training pass
    optimizer.zero_grad()

    output = model.forward(images)

    loss = criterion(output, labels)
    loss.backward()
    optimizer.step()

My model worked without any problems when using cross entropy loss, but when I changed to MSE loss, the interpreter complained and said that my tensors were different sizes and thus could not be computed.

<class 'torch.Tensor'>
torch.Size([64, 1, 28, 28])
torch.Size([64])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-62-ec6942122f02> in <module>
     44     output = model.forward(images)
     45 
---> 46     loss = criterion(output, labels)
     47     loss.backward()
     48     optimizer.step()

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    530             result = self._slow_forward(*input, **kwargs)
    531         else:
--> 532             result = self.forward(*input, **kwargs)
    533         for hook in self._forward_hooks.values():
    534             hook_result = hook(self, input, result)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    429 
    430     def forward(self, input, target):
--> 431         return F.mse_loss(input, target, reduction=self.reduction)
    432 
    433 

/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py in mse_loss(input, target, size_average, reduce, reduction)
   2213             ret = torch.mean(ret) if reduction == 'mean' else torch.sum(ret)
   2214     else:
-> 2215         expanded_input, expanded_target = torch.broadcast_tensors(input, target)
   2216         ret = torch._C._nn.mse_loss(expanded_input, expanded_target, _Reduction.get_enum(reduction))
   2217     return ret

/opt/conda/lib/python3.7/site-packages/torch/functional.py in broadcast_tensors(*tensors)
     50                 [0, 1, 2]])
     51     """
---> 52     return torch._C._VariableFunctions.broadcast_tensors(tensors)
     53 
     54 

RuntimeError: The size of tensor a (10) must match the size of tensor b (64) at non-singleton dimension 1

I tried reshaping my tensors and creating new arrays as placeholders for my output array, yet seem to be getting nowhere.

Why cross entropy loss works without any errors yet MSE does not?

1 Answers

nn.CrossEntropyLoss and nn.MSELoss are completely different loss functions with fundamentally different rationale behind them.

nn.CrossEntropyLoss is a loss function for discrete labeling tasks. Therefore it expects as inputs a prediction of label probabilities and targets as ground-truth discrete labels: x shape is nxc (where c is the number of labels) and y is of shape n of type integer, each target takes values in the range {0,...,c-1}.

In contrast, nn.MSELoss is a loss function for regression tasks. Therefore it expects both predictions and targets to be of the same shape and data type. That is, if your prediction is of shape nxc the target should also be of shape nxc (and not just n as in the cross-entropy case).

If you are insisting on using MSE loss instead of cross entropy, you will need to convert the target integer labels you currently have (of shape n) into 1-hot vectors of shape nxc and only then compute the MSE loss between your predictions and the generated one-hot targets.

Related