用 PyTorch 实现损失函数：交叉熵、均方误差与自定义函数

损失函数把模型输出与训练目标之间的差异变成一个可优化的数值。选择合适的损失函数，不只是挑一个 API：还要确认模型输出代表什么、目标张量采用什么格式、各维度如何归约，以及哪些样本应该参与计算。形状或语义不匹配时，代码有时仍能运行，却会优化错误的目标。

本文介绍 PyTorch 中常用的交叉熵和均方误差，说明它们适用的任务与参数边界，并从零实现带掩码的自定义损失。示例使用 CPU 和随机张量，可直接运行。示例依赖 PyTorch 2.x 常用 API；其中所用的基本接口在较早版本中也很常见，具体参数支持情况应以当前安装版本的文档为准。

一、先检查三件事：输出、目标与归约

写损失函数前，先回答三个问题：

• 模型输出是什么？是未归一化的分类分数（logits）、概率，还是连续数值？
• 目标如何表示？是类别编号、类别概率，还是与预测同形状的实数？
• 损失怎样汇总？对所有元素取平均、求和，还是保留每个元素的损失？

下面的代码统一在训练前检查形状和数据类型。对于使用自动混合精度的训练，损失计算一般也能参与 PyTorch 的类型转换流程；若遇到精度或溢出问题，应结合训练方式检查实际张量的 dtype，而不是默认所有输入都为 float32。

二、交叉熵：分类任务的常用选择

torch.nn.CrossEntropyLoss适用于互斥类别的分类任务，例如一张图片属于猫、狗或鸟中的一个类别。该接口接收未经过 softmax 的 logits，内部会处理归一化和负对数似然计算。不要先对输出调用 softmax 再传入，否则会改变损失所依据的数值，通常不是想要的计算方式。

2.1 单标签分类：目标是类别编号

分类头输出形状通常为 (N, C)：N 是批量大小，C 是类别数。目标为形状 (N,) 的整数张量，每个值表示正确类别的索引，通常使用 torch.long。类别索引必须处于 [0, C)；如果设置了 ignore_index，该指定值可以作为例外。

import torch
import torch.nn as nn

torch.manual_seed(0)

batch_size = 4
num_classes = 3

logits = torch.randn(batch_size, num_classes)  # 未归一化分类分数
target = torch.tensor([0, 2, 1, 0], dtype=torch.long)

criterion = nn.CrossEntropyLoss()
loss = criterion(logits, target)

print("logits:", logits.shape)
print("target:", target.shape, target.dtype)
print("loss:", loss.item())

若分类目标本身是概率分布，也可以传入与 logits 同形状的浮点目标，例如 (N, C)。这种模式适用于软标签等场景。需要特别注意，PyTorch 不会替使用者验证每行概率非负且总和为 1；传入不符合预期的数值可能不会报错，但损失的含义也可能不再符合概率分布下的交叉熵。

2.2 分割和序列分类：多出空间或时间维

交叉熵也可用于逐像素分割或逐位置分类。输入形状为 (N, C, d1, d2, ...)，类别维固定在第 2 维；使用类别索引作为目标时，目标形状为 (N, d1, d2, ...)。例如输入 (N, C, H, W)，目标应为 (N, H, W)，而不是再带一个类别维。

下面演示二维分割中的忽略标签。ignore_index用于让指定目标值不参与损失计算，适合标注中存在“未标注”区域的情况。它不会自动修复越界标签：除该忽略值外，类别索引仍应落在有效范围内。

import torch
import torch.nn as nn

batch_size, num_classes, height, width = 2, 4, 5, 6
logits = torch.randn(batch_size, num_classes, height, width)
target = torch.randint(0, num_classes, (batch_size, height, width))

ignore_index = 255
target[0, 0, 0] = ignore_index

criterion = nn.CrossEntropyLoss(ignore_index=ignore_index)
loss = criterion(logits, target)

print("loss:", loss.item())

2.3 常用参数与边界

• weight：可选的类别权重，通常是长度为 C 的一维张量，用于调整各类别对损失的贡献。它应与类别数、设备和计算类型相匹配。类别权重不等同于样本权重。
• ignore_index：在类别索引目标模式下忽略指定标签。它不能让任意越界标签都变得有效，也不是概率目标模式下用于屏蔽位置的通用参数。
• reduction：可取 "none"、"mean" 或 "sum"。"none"保留逐样本或逐位置结果；"sum"求和；"mean"求平均。使用类别权重、忽略标签时，平均的具体分母会受目标和权重影响，不应简单理解成始终除以批量大小。
• label_smoothing：用于平滑类别目标，参数应在 [0, 1] 范围内。它不是额外的类别，也不能替代对标签格式和模型输出的检查。

交叉熵最常见的错误，是把二分类、多标签分类和多类别单标签分类混为一谈。若每个样本恰好属于多个互斥类别中的一个，通常使用交叉熵；若每个样本可以同时具有多个独立标签，通常考虑逐标签的二元交叉熵，例如 BCEWithLogitsLoss。后者接收 logits 和与 logits 对齐的浮点目标，目标值通常表示 0 到 1 之间的标签值。不要对 logits 先做 sigmoid 再传给 BCEWithLogitsLoss；该接口已将相应计算组合在一起。

三、均方误差：回归与数值拟合

torch.nn.MSELoss计算预测值与目标值逐元素差值的平方。若逐元素差值为 e，则每个元素的误差为 e²。它适用于预测连续数值的回归任务，也常用于重建、蒸馏等对齐数值张量的任务。

默认的 reduction="mean" 会对所有元素的平方误差取平均，而不只是对批量维求平均。因此，输出有多个特征时，损失是跨批次和特征维度共同归约的。若需要不同的汇总逻辑，可以先设置 reduction="none"，再自行聚合。

import torch
import torch.nn as nn

prediction = torch.tensor(
    [[2.0, 1.0], [0.5, 3.0]],
    requires_grad=True
)
target = torch.tensor(
    [[1.0, 1.0], [1.5, 2.0]]
)

criterion = nn.MSELoss()
loss = criterion(prediction, target)
loss.backward()

print("loss:", loss.item())
print("gradient:", prediction.grad)

使用均方误差时，预测和目标通常应具有相同形状，并表示可比较的数值。如果它们形状不同，PyTorch 有时会按广播规则计算，同时发出警告；这不代表结果符合预期。比如 (N, 1) 与 (N,) 可能广播成 (N, N)，使每个预测都与整批目标组合计算。更稳妥的做法是在进入损失函数前断言形状相同。

import torch
import torch.nn as nn

prediction = torch.randn(8, 1)
target = torch.randn(8)

# 目标补上特征维，避免将 (8, 1) 与 (8,) 广播为 (8, 8)
target = target.unsqueeze(1)

assert prediction.shape == target.shape
loss = nn.MSELoss()(prediction, target)
print(loss.item())

MSE 会对较大的误差平方放大，因此离群目标可能对训练产生较大影响。如果数据中存在异常值，是否改用其他回归损失应根据任务目标和数据分布决定，不能仅凭“更稳健”就机械替换。还要检查预测和目标的尺度、归一化方式以及缺失值处理；含有 NaN 的输入通常会让损失结果也变成 NaN。

四、从公式实现自定义损失

当内置损失的目标、掩码或归约方式不符合需求时，可以直接用 PyTorch 张量运算实现。只要运算仍在张量计算图中，梯度通常可以自动传播。下面实现带掩码的逐元素加权平方误差：掩码为 1 的位置参与计算，为 0 的位置忽略；权重可以按可广播的形状提供。

这里选择按有效权重之和归一化，而不是对整张张量取平均。这样屏蔽区域不会因为数量变化而改变分母。若所有位置都被屏蔽，分母为 0；代码明确返回连接计算图的零值，避免除零，同时调用方仍应决定这种批次是否合理。

import torch

def masked_weighted_mse(prediction, target, mask, weight=None):
    """
    prediction、target：形状相同的浮点张量
    mask：可广播到 prediction 形状的 0/1 张量
    weight：可选的非负权重，可广播到 prediction 形状
    """
    if prediction.shape != target.shape:
        raise ValueError(
            f"prediction 和 target 形状必须相同："
            f"{prediction.shape} != {target.shape}"
        )

    if not prediction.is_floating_point():
        raise TypeError("prediction 必须是浮点张量")
    if not target.is_floating_point():
        raise TypeError("target 必须是浮点张量")

    mask = mask.to(device=prediction.device, dtype=prediction.dtype)

    if weight is None:
        weight = torch.ones_like(prediction)
    else:
        weight = weight.to(device=prediction.device, dtype=prediction.dtype)

    # 先由张量运算自然检查 mask、weight 是否能广播到预测形状
    effective_weight = mask * weight
    squared_error = (prediction - target).square()
    numerator = (squared_error * effective_weight).sum()
    denominator = effective_weight.sum()

    # 分母为零时返回与 prediction 保持计算图连接的零值
    safe_denominator = denominator.clamp_min(
        torch.finfo(prediction.dtype).tiny
    )
    result = numerator / safe_denominator
    return torch.where(
        denominator > 0,
        result,
        prediction.sum() * 0.0
    )


prediction = torch.tensor(
    [[1.0, 3.0], [2.0, 5.0]],
    requires_grad=True
)
target = torch.tensor([[0.0, 2.0], [3.0, 4.0]])
mask = torch.tensor([[1.0, 0.0], [1.0, 1.0]])
weight = torch.tensor([[1.0, 1.0], [2.0, 1.0]])

loss = masked_weighted_mse(prediction, target, mask, weight)
loss.backward()

print("loss:", loss.item())
print("gradient:", prediction.grad)

这个实现假定掩码和权重能够广播到预测形状，且权重非负。函数没有逐项扫描张量并验证掩码只包含 0 和 1、权重没有负数或 NaN；若输入由外部数据生成，应在数据入口处按应用需要校验这些条件。将张量转到预测所在设备和类型，能够避免常见的 CPU/GPU 混用错误，但不会消除所有精度或数值稳定性问题。

自定义函数还可以返回逐元素损失，留给训练循环决定归约方式。例如：

import torch

def elementwise_huber_like(prediction, target, delta=1.0):
    """演示用分段损失；要求 delta 为正数。"""
    if prediction.shape != target.shape:
        raise ValueError("prediction 和 target 形状必须相同")
    if delta <= 0:
        raise ValueError("delta 必须大于 0")

    error = (prediction - target).abs()
    quadratic = torch.minimum(error, torch.as_tensor(
        delta, device=error.device, dtype=error.dtype
    ))
    linear = error - quadratic
    return 0.5 * quadratic.square() + delta * linear


prediction = torch.tensor([0.2, 2.0], requires_grad=True)
target = torch.tensor([0.0, 0.0])
per_element = elementwise_huber_like(prediction, target, delta=1.0)
loss = per_element.mean()
loss.backward()

print("逐元素损失:", per_element)
print("平均损失:", loss.item())

实际项目中，如果自定义目标能由内置损失和权重、掩码等组合清晰表达，优先复用内置实现通常更容易维护。自定义损失应写明公式、形状约定、归约方式、有效输入范围，并用小张量手算验证前向值；需要训练时，还应确认梯度没有意外截断。

五、常见错误检查清单

• 交叉熵前先做 softmax。对 CrossEntropyLoss传入原始 logits；概率目标模式是另一种明确的数据格式，不等于“把预测先 softmax 后仍当 logits 使用”。
• 类别编号用错 dtype 或形状。索引目标通常应为 torch.long，形状不包含类别维；分割目标的空间维应与输入对应。
• 类别索引越界。索引应小于类别数，除非它恰好等于配置的 ignore_index。软目标则应与输入同形状，并由调用方保证数值语义正确。
• 回归目标被意外广播。MSE 前显式检查预测和目标形状，尤其留意末尾维度为 1 的情况。
• 把“平均”当成按批量大小平均。MSE 的默认平均覆盖所有元素；自定义加权损失的分母也必须按预定公式选择。
• 标签和预测落在不同设备。损失计算所需张量应处于兼容设备。训练中发现设备错误时，检查数据加载、目标构造和模型输出的设备。
• 损失出现 NaN 或无穷值。检查输入是否已含非有限值、分母是否为零、权重范围、数据尺度及自定义公式的数值边界。只用 nan_to_num掩盖异常，可能让训练继续优化错误数据。

六、参数与安全边界

损失函数的参数限制是计算接口的约束，不是对数据正确性的完整保证。例如，交叉熵的 label_smoothing应在 [0, 1]，类别权重应与类别数对齐；MSE 的输入需要具有可比较的数值含义；自定义掩码和权重则应符合代码设定的形状与取值约定。对来自文件、网络或用户的配置，建议在进入训练前显式校验形状、类型、有限值和取值范围，并为失败情况给出清楚的错误信息。

损失函数不是安全机制：它不会验证标签是否可信，不会限制模型输出可访问的资源，也不能证明模型对抗攻击、数据投毒或越权输入是安全的。不要把较低的验证损失等同于安全保证。若训练或推理流程会读取外部提供的模型文件，应采用适当的来源校验和文件访问控制；不要对不可信文件使用可能执行任意代码的反序列化方式。损失函数的选择和安全评估属于不同问题，需分别处理。

总结

互斥类别分类通常从交叉熵开始，记住输入是 logits，并按任务选择类别索引或软目标格式；连续数值拟合可使用 MSE，同时严格检查预测和目标形状，避免意外广播；当需要掩码、特殊归约或自定义目标时，先明确公式，再用张量运算实现并验证边界。写对损失函数，关键不在于 API 数量，而在于输出与目标的语义、维度和归约规则彼此一致。