PyTorch中的损失函数：F.nll_loss 与 nn.CrossEntropyLoss

背景介绍

无论是图像分类、文本分类还是其他类型的分类任务，交叉熵损失（Cross Entropy Loss）都是最常用的一种损失函数。它衡量的是模型预测的概率分布与真实标签之间的差异。在 PyTorch 中，有两个特别值得注意的实现：F.nll_loss 和 nn.CrossEntropyLoss。

F.nll_loss

什么是负对数似然损失？

F.nll_loss 是负对数似然损失（Negative Log Likelihood Loss），主要用于多类分类问题。它的输入是对数概率（log-probabilities），这意味着在使用 F.nll_loss 之前，我们需要先对模型的输出应用 log_softmax 函数，将原始输出转换为对数概率形式。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

# 创建一些虚拟数据
features = torch.randn(100, 20)  # 假设有100个样本，每个样本有20个特征
labels = torch.randint(0, 3, (100,))  # 假设有3个类别

# 创建数据加载器
dataset = TensorDataset(features, labels)
data_loader = DataLoader(dataset, batch_size=10, shuffle=True)

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(20, 3)  # 输入维度为20，输出维度为3（对应3个类别）

    def forward(self, x):
        return self.fc(x)

model_nll = SimpleModel()
optimizer = torch.optim.SGD(model_nll.parameters(), lr=0.01)

for inputs, targets in data_loader:
    optimizer.zero_grad()  # 清除梯度
    outputs = model_nll(inputs)  # 模型前向传播
    log_softmax_outputs = F.log_softmax(outputs, dim=1)  # 应用 log_softmax
    loss = F.nll_loss(log_softmax_outputs, targets)  # 计算 nll_loss
    loss.backward()  # 反向传播
    optimizer.step()  # 更新权重

    print(f"Batch Loss with F.nll_loss: {loss.item():.4f}")

应用场景

由于 F.nll_loss 需要预先计算 log_softmax，这为用户提供了一定程度的灵活性，尤其是在需要复用 log_softmax 结果的情况下。

nn.CrossEntropyLoss

简化工作流程

相比之下，nn.CrossEntropyLoss 更加直接和易用。它结合了 log_softmax 和 nll_loss 的功能，因此可以直接接受未经归一化的原始输出作为输入，内部自动完成这两个步骤。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

# 创建一些虚拟数据
features = torch.randn(100, 20)  # 假设有100个样本，每个样本有20个特征
labels = torch.randint(0, 3, (100,))  # 假设有3个类别

# 创建数据加载器
dataset = TensorDataset(features, labels)
data_loader = DataLoader(dataset, batch_size=10, shuffle=True)

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(20, 3)  # 输入维度为20，输出维度为3（对应3个类别）

    def forward(self, x):
        return self.fc(x)

model_ce = SimpleModel()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model_ce.parameters(), lr=0.01)

for inputs, targets in data_loader:
    optimizer.zero_grad()  # 清除梯度
    outputs = model_ce(inputs)  # 模型前向传播
    loss = criterion(outputs, targets)  # 直接计算交叉熵损失，内部包含 log_softmax
    loss.backward()  # 反向传播
    optimizer.step()  # 更新权重

    print(f"Batch Loss with nn.CrossEntropyLoss: {loss.item():.4f}")

内部机制

实际上，nn.CrossEntropyLoss = log_softmax + nll_loss 。这种设计简化了用户的代码编写过程，特别是当不需要对中间结果进行额外操作时。

区别与联系

• 输入要求：F.nll_loss 要求输入为 log_softmax 后的结果；而 nn.CrossEntropyLoss 可以直接接受未经 softmax 处理的原始输出。

• 灵活性：如果需要对 log_softmax 结果进行进一步处理或调试，那么 F.nll_loss 提供了更大的灵活性。

• 便捷性：对于大多数用户而言，nn.CrossEntropyLoss 因其简洁性和内置的 log_softmax 步骤，是更方便的选择。