What is "cross entropy loss" really doing when input is 3D?

Na1ve 0
Viewed 20

I'm working on a text-generating RNN. I found out that when calculating cross entropy loss, if the input has a size of[batch_size, vocab_size, seq_len] and the target has a size of[batch_size, seq_len], then the model is not working at all however long I train it.

But, if I resize the input to [batch_size*seq_len, vocab] and the target to a 1D vector, then the model will be working just fine.

So my question is, what is cross entropy loss really doing when handling 3D input? Why the first type of calculating loss won't work in my task?

cross-entropy loss-function text pytorch
1 Answers

They should be the same. Something else is wrong and it is insufficient to deduce from your question.

Code snippet showing that both the loss values and the gradients are the same:

import torch
import torch.nn as nn

loss = nn.CrossEntropyLoss()

batch_size = 4
seq_len = 8
vocab_size = 3

inputs = torch.randn((batch_size, vocab_size, seq_len), requires_grad=True)
target = torch.randint(low=0, high=vocab_size, size=(batch_size, seq_len))

loss1 = loss(inputs, target)
grad1 = torch.autograd.grad(loss1, inputs)[0]

inputs_transposed = inputs.permute(0, 2, 1).reshape(batch_size*seq_len, vocab_size)
target_transposed = target.view(batch_size*seq_len)

loss2 = loss(inputs_transposed, target_transposed)
grad2 = torch.autograd.grad(loss2, inputs)[0]

print(torch.allclose(loss1, loss2))
print(torch.allclose(grad1, grad2))
hkchengrex0
Related