PyTorch自定义Loss必须继承nn.Module或torch.autograd.Function才能参与反向传播；普通函数因不构建计算图而无法求导；继承nn.Module需在forward中实现可导运算并调用super().__init__()；需避免.item()、.numpy()等中断梯度操作；复杂场景（如不可导操作、梯度定制）应使用autograd.Function，显式定义forward与backward；务必验证loss.grad_fn非None且参数梯度不为None。
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的输出/输入一一对应。
立即学习“Python免费学习笔记（深入）”；
例如实现一个带梯度缩放的L1 Loss：
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）
多输出Loss或带辅助项时，别忘了梯度归一化
比如在目标Loss外加一个正则项：total_loss = task_loss + 0.01 * reg_loss。如果reg_loss量级远大于task_loss（如L2 norm比交叉熵大几个数量级），梯度会严重偏向正则项，训练失衡。
- 建议对每项Loss单独
.mean()或.sum()后再加权，而不是直接加未归一化的张量 - 调试时可分别打印
task_loss.item()和reg_loss.item()，观察量级是否合理 - 若用
autograd.Function实现复合Loss，backward里返回的梯度也要按相同系数缩放，否则梯度更新步长错乱
最难察觉的是Loss里隐式改变了tensor的requires_grad状态，比如通过torch.where选值后忘记保留梯度路径，或者用torch.max取index再索引——这些操作容易切断计算图。动手前先用torch.autograd.gradcheck做数值梯度校验。
