torch.nn.functional.cross_entropy
是PyTorch中用于计算交叉熵损失的函数,非常适合用于多分类问题。这个函数结合了log_softmax
操作和nll_loss
(负对数似然损失)的计算,因此输入得分(即模型输出)不需要事先经过softmax处理。
下面是一个使用torch.nn.functional.cross_entropy
的示例,展示了如何在一个简单的神经网络模型中应用它来计算损失:
python
import torch
import torch.nn as nn
import torch.nn.functional as F
# 假设我们有一个简单的模型
class SimpleModel(nn.Module):
def __init__(self, input_size, num_classes):
super(SimpleModel, self).__init__()
self.linear = nn.Linear(input_size, num_classes)
def forward(self, x):
return self.linear(x)
# 创建模型实例
model = SimpleModel(input_size=10, num_classes=3)
# 模拟一批输入数据和标签
inputs = torch.randn(5, 10) # 假设批大小为5,输入特征大小为10
labels = torch.tensor([0, 2, 1, 0, 2]) # 真实标签
# 模型前向传播
outputs = model(inputs)
# 计算交叉熵损失
# outputs: torch.Size([5, 3])
# labels: torch.Size([5]), 注意这个参数必须为long型的
loss = F.cross_entropy(outputs, labels)
print("Loss:", loss.item())
在这个示例中:
SimpleModel
是一个简单的线性模型,其输出大小等于类别数。- 我们创建了一批输入
inputs
和对应的标签labels
。 - 模型的输出
outputs
是直接传递给F.cross_entropy
的,不需要额外的softmax层,因为cross_entropy
内部已经处理了这部分。 labels
应该是每个样本的类别索引形式,而不是one-hot编码。F.cross_entropy
计算了从模型输出到真实标签的交叉熵损失。
这种方式使得实现多分类问题的模型训练变得简单而直接。