【pytorch】nn.CrossEntropyLoss()

nn.CrossEntropyLoss()函数计算交叉熵损失

import torch

# -----------------------------------------
# 定义数据: batch_size=4；  一共有5个分类
# label.size() : torch.Size([4])
# predict.size(): torch.Size([4, 5])
# -----------------------------------------
torch.manual_seed(100)
label = torch.tensor([4, 3, 3, 2])
predict = torch.rand(4, 5)


# # -----------------------------------------
# # 直接调用函数 nn.CrossEntropyLoss() 计算 Loss
# # -----------------------------------------
# criterion = nn.CrossEntropyLoss()
# loss = criterion(predict, label)
# print(loss)   # tensor(1.5230)


# -----------------------------------------
# 自己根据 公式手撸代码
# -----------------------------------------
def cross_entropy_loss(predict, label):
    temp1 = torch.gather(predict, dim=1, index=label.view(-1, 1))
    temp2 = torch.exp(temp1) / torch.sum(torch.exp(predict), dim=1, keepdim=True)
    loss = -torch.log(temp2)
    return torch.sum(loss) / len(predict)


loss = cross_entropy_loss(predict, label)
print(loss)   # tensor(1.5230)