[FIXED] pytorch nllloss function target shape mismatch

Issue

I'm training a LSTM model using pytorch with batch size of 256 and NLLLoss() as loss function. The loss function is having problem with the data shape.

The softmax output from the forward passing has shape of torch.Size([256, 4, 1181]) where 256 is batch size, 4 is sequence length, and 1181 is vocab size.

The target is in the shape of torch.Size([256, 4]) where 256 is batch size and 4 is the output sequence length.

When I was testing earlier with batch size of 1, the model works fine but when I add batch size, it is breaking. I read that NLLLoss() can take class target as input instead of one hot encoded target.

Am I misunderstanding it? Or did I not format the shape of the target correctly?

class LSTM(nn.Module):

    def __init__(self, embed_size=100, hidden_size=100, vocab_size=1181, embedding_matrix=...):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.word_embeddings = nn.Embedding(vocab_size, embed_size)
        self.word_embeddings.load_state_dict({'weight': torch.Tensor(embedding_matrix)})
        self.word_embeddings.weight.requires_grad = False
        self.lstm = nn.LSTM(embed_size, hidden_size)
        self.hidden2out = nn.Linear(hidden_size, vocab_size)


    def forward(self, tokens):
        batch_size, num_steps = tokens.shape
        embeds = self.word_embeddings(tokens)
        lstm_out, _ = self.lstm(embeds.view(batch_size, num_steps, -1))
        out_space = self.hidden2out(lstm_out.view(batch_size, num_steps, -1))
        out_scores = F.log_softmax(out_space, dim=1)
        return out_scores

model = LSTM(self.config.embed_size, self.config.hidden_size, self.config.vocab_size, self.embedding_matrix)
loss_function = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=self.config.lr)

Error:

~/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   1846         if target.size()[1:] != input.size()[2:]:
   1847             raise ValueError('Expected target size {}, got {}'.format(
-> 1848                 out_size, target.size()))
   1849         input = input.contiguous().view(n, c, 1, -1)
   1850         target = target.contiguous().view(n, 1, -1)

ValueError: Expected target size (256, 554), got torch.Size([256, 4])

Solution

Your input shape to the loss function is (N, d, C) = (256, 4, 1181) and your target shape is (N, d) = (256, 4), however, according to the docs on NLLLoss the input should be (N, C, d) for a target of (N, d).

Supposing x is your network output and y is the target then you can compute loss by transposing the incorrect dimensions of x as follows:

loss = loss_function(x.transpose(1, 2), y)

Alternatively, since NLLLoss is just averaging all the responses anyway, you can reshape x and y to be (N*d, C) and (N*d). This gives the same result without creating temporary copies of your tensors.

loss = loss_function(x.reshape(N*d, C), y.reshape(N*d))

Answered By - jodag