如何在Python中自定义PyTorch的Loss损失函数与反向传播？-Python教程-PHP中文网

PyTorch自定义Loss必须继承nn.Module或torch.autograd.Function才能参与反向传播；普通函数因不构建计算图而无法求导；继承nn.Module需在forward中实现可导运算并调用super().__init__()；需避免.item()、.numpy()等中断梯度操作；复杂场景（如不可导操作、梯度定制）应使用autograd.Function，显式定义forward与backward；务必验证loss.grad_fn非None且参数梯度不为None。

如何在python中自定义pytorch的loss损失函数与反向传播？

PyTorch自定义Loss必须继承`torch.nn.Module`或重写`forward`

直接写一个普通函数（比如def my_loss(y_pred, y_true): return torch.mean((y_pred - y_true) ** 2)）虽然能算出数值，但无法参与反向传播——因为没构建计算图。PyTorch的自动求导只追踪nn.Module子类或torch.autograd.Function定义的前向/反向逻辑。

最稳妥的做法是继承nn.Module，在forward里实现计算，并确保所有张量操作都是可导的：

class MyMSELoss(torch.nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction
<pre class="brush:php;toolbar:false;">def forward(self, input, target):
    loss = (input - target) ** 2
    if self.reduction == 'mean':
        return loss.mean()
    elif self.reduction == 'sum':
        return loss.sum()
    else:
        return loss

别漏掉super().__init__()，否则self.reduction不会被注册为模块属性
reduction参数要自己实现，PyTorch不会自动帮你处理
避免在forward里用.item()、.numpy()或print()等中断梯度的操作

需要控制反向传播细节时，用`torch.autograd.Function`

当Loss里包含不可导操作（如argmax）、想手动覆写梯度、或需节省显存（比如梯度裁剪嵌入Loss），就得用autograd.Function。它要求显式定义forward和backward，且backward的输入输出必须与forward的输出/输入一一对应。

例如实现一个带梯度缩放的L1 Loss：

python-code-analyz

专业Python代码分析与优化，支持语法检查、安全扫描、性能评估、复杂度分析及重构后优化代码生成。

下载

class ScaledL1Loss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, target, scale=1.0):
        ctx.save_for_backward(input, target)
        ctx.scale = scale
        return torch.mean(torch.abs(input - target))
<pre class="brush:php;toolbar:false;">@staticmethod
def backward(ctx, grad_output):
    input, target = ctx.saved_tensors
    grad_input = grad_output * ctx.scale * torch.sign(input - target) / input.numel()
    return grad_input, None, None  # input有梯度，target和scale没有

ctx.save_for_backward()保存中间变量供backward使用；不能直接存Python标量（如scale），要存张量或用ctx.xxx存
backward返回的梯度顺序必须和forward参数顺序一致；不需要梯度的参数返回None
调用时要用ScaledL1Loss.apply(input, target, 0.1)，不是ScaledL1Loss()(…)

验证自定义Loss是否真正参与反向传播

光看loss值下降不够，得确认梯度确实回传到了模型参数。常见错误是Loss返回了标量但没连上模型输出，或用了.detach()、with torch.no_grad():包裹。

检查Loss输出是否有grad_fn：loss.grad_fn不为None才说明在计算图中
运行后立刻打印某层权重的梯度：model.layer.weight.grad应为非None的Tensor
如果loss.backward()后所有grad都是None，大概率是Loss里用了.item()、numpy()，或输入input本身requires_grad=False
注意input和target的device和dtype要一致，否则可能静默失败（如float32 vs float64）