如何在Python中实现PyTorch自定义损失函数_继承_Loss类或直接编写函数-Python教程-PHP中文网

必须继承 nn.Module，不能只写纯函数；因训练循环依赖 loss.backward()，仅 nn.Module 实例返回的 Tensor 才能保留计算图，普通函数易致梯度静默失败或 RuntimeError。

如何在python中实现pytorch自定义损失函数_继承_loss类或直接编写函数

PyTorch自定义损失函数：该继承 `nn.Module` 还是写纯函数？

直接说结论：**必须继承 nn.Module（或其子类 nn.Loss）**，不能只写一个普通 Python 函数。PyTorch 的训练循环（如 optimizer.step()）依赖 loss.backward() 触发自动求导，而只有 nn.Module 实例返回的 Tensor 才能正确保留计算图——普通函数若未显式确保 requires_grad=True 且参与图构建，反向传播会静默失败或报 RuntimeError: element 0 of tensors does not require grad。

继承 `nn.Module` 实现损失类的最小必要结构

核心是重写 forward 方法，并确保输入 pred 和 target 都参与计算图。不需要重写 __init__（除非要传超参），也不用调用 super().__init__() ——但建议加上，便于未来扩展。

常见错误：

在 forward 中用了 numpy 或 .item() / .detach().cpu().numpy()，导致计算图断裂
返回标量时用了 torch.mean(loss).item()，返回的是 Python float，不是 Tensor
对 target 做了非可导操作（如 torch.argmax 后再算交叉熵），应改用原生支持的 nn.CrossEntropyLoss 或手动实现 soft 版本

正确示例（带权重的二分类 focal loss）：

import torch
import torch.nn as nn
<p>class FocalLoss(nn.Module):
def <strong>init</strong>(self, alpha=1.0, gamma=2.0):
super().<strong>init</strong>()
self.alpha = alpha
self.gamma = gamma</p><pre class="brush:php;toolbar:false;">def forward(self, inputs, targets):
    # inputs: [N, 2], targets: [N] with class indices (0 or 1)
    logpt = nn.functional.log_softmax(inputs, dim=1)
    pt = torch.exp(logpt)
    logpt = logpt.gather(1, targets.unsqueeze(1))
    pt = pt.gather(1, targets.unsqueeze(1))
    focal_weight = ((1 - pt) ** self.gamma)
    loss = -self.alpha * focal_weight * logpt
    return loss.mean()

python 查询技能

查询客流数据，输出JSON格式，可直接导入Bitable等可视化工具

下载

为什么不用 `nn.Loss` 基类？

PyTorch 并没有名为 nn.Loss 的抽象基类。官方所有损失函数（如 nn.MSELoss、nn.CrossEntropyLoss）都直接继承自 nn.Module。所谓“继承 Loss 类”是误传——你看到的文档里写的 “Base class for all losses” 实际指向的是内部的 _Loss（带下划线，非公开），用户不应继承它。继承 nn.Module 是唯一受支持、稳定且可维护的方式。

使用场景差异：