Cross Entropy

cross_entropy_1d_2d.py

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F

num_classes = 5
input = torch.randn(3, num_classes, requires_grad=True)
print(input)
target = torch.randint(num_classes, (3,), dtype=torch.int64)
print(target)

softmax = nn.Softmax(dim=1)
softmaxed = softmax(input)
print(softmaxed)
loss = F.cross_entropy(softmaxed, target, reduction='none')
print(loss)

input2d = torch.randn(3, num_classes, 7, 7, requires_grad=True)
target2d = torch.randint(num_classes, (3, 7, 7), dtype=torch.int64)
print(input2d)
print(target2d)
softmax = nn.Softmax(dim=2)
softmaxed = softmax(input2d)

loss = F.cross_entropy(softmaxed, target2d)
print(loss)

Note that
input2d.size() -> torch.Size([3, 5, 7, 7])
target2d.size() -> torch.Size([3, 7, 7])