Pytorch crossentropy loss with 3d input

user_6077443 0
Viewed 2292

I have a network which outputs a 3D tensor of size (batch_size, max_len, num_classes). My groud truth is in the shape (batch_size, max_len). If I do perform one-hot encoding on the labels, it'll be of shape (batch_size, max_len, num_classes) i.e the values in max_len are integers in the range [0, num_classes]. Since the original code is too long, I have written a simpler version that reproduces the original error.

criterion = nn.CrossEntropyLoss()
batch_size = 32
max_len = 350
num_classes = 1000
pred = torch.randn([batch_size, max_len, num_classes])
label = torch.randint(0, num_classes,[batch_size, max_len])
pred = nn.Softmax(dim = 2)(pred)
criterion(pred, label)

the shape of pred and label are respectively,torch.Size([32, 350, 1000]) and torch.Size([32, 350])

The error encountered is

ValueError: Expected target size (32, 1000), got torch.Size([32, 350, 1000])

If I one-hot encode labels for computing the loss

x = nn.functional.one_hot(label)
criterion(pred, x)

it'll throw the following error

ValueError: Expected target size (32, 1000), got torch.Size([32, 350, 1000])
cross-entropy pytorch neural-network python
1 Answers

From the Pytorch documentation, CrossEntropyLoss expects the shape of its input to be (N, C, ...), so the second dimension is always the number of classes. Your code should work if you reshape preds to be of size (batch_size, num_classes, max_len).

user_124189530
Related