PyTorch中多任务损失必须手动加权求和，因无内置封装；需分别计算各任务loss并加权组合，注意量纲归一化、梯度路由验证及独立指标评估。

在python中如何通过pytorch实现多任务学习的自定义损失函数？

PyTorch中多任务损失函数必须手动加权求和

PyTorch本身不提供开箱即用的多任务损失封装，nn.Module 的子类必须显式定义各任务损失的计算路径和组合逻辑。常见错误是直接把多个 loss1 和 loss2 相加却不加权重，导致梯度被某一任务主导——比如分类任务 loss 值常为 0.5～2.0，而回归任务 MSE 可能高达 100+，不缩放会彻底压垮小量级任务的更新。

实操建议：

每个任务单独调用对应 loss 函数（如 nn.CrossEntropyLoss()、nn.MSELoss()），不要复用同一个实例去算不同输出
用可学习或固定权重控制平衡：total_loss = w1 * loss_cls + w2 * loss_reg，其中 w1、w2 初始设为 1.0，后续可按验证集各任务性能动态调整
若任务量纲差异极大（如像素值回归 vs 类别预测），先对回归目标做归一化（如除以训练集标准差），再算 loss，比硬调权重更稳定

共享主干网络时，梯度反向传播不会自动隔离

多任务模型常用共享 backbone + 多个 task head 结构，但 PyTorch 默认反向传播会把所有任务 loss 的梯度叠加到共享参数上。这看似合理，实际容易引发梯度冲突——比如一个任务希望某层卷积核增强纹理响应，另一个任务却要求抑制该响应。

实操建议：

检查 model.shared_backbone 的 requires_grad 是否为 True（应为 True），同时确认各 task head 的参数也未被意外冻结
避免在 forward 中用 torch.no_grad() 包裹某一分支——它会切断整个子图的梯度流，导致该任务 loss 对 backbone 完全无影响
调试时打印梯度 norm：print(torch.norm(model.shared_backbone[0].weight.grad).item())，若某任务 loss 设为 0 后该 norm 不变，说明梯度未正确路由

自定义 loss 类需重写 call 而非 forward

很多人照搬模型写法，在自定义 loss 类里写 forward() 方法，结果发现无法像 nn.CrossEntropyLoss() 那样直接调用。PyTorch 的 loss 模块本质是 callable 对象，核心是实现 __call__，内部再委托给 forward 或直接计算。

实操示例：

Shadows Python Sensei

Python 最佳实践助手——代码规范、设计模式、性能优化、测试与类型注解。适用于编写或审查 Python 代码。

下载

class MultiTaskLoss(nn.Module):
    def __init__(self, w_cls=1.0, w_reg=1.0):
        super().__init__()
        self.w_cls = w_cls
        self.w_reg = w_reg
        self.cls_loss = nn.CrossEntropyLoss()
        self.reg_loss = nn.MSELoss()
<pre class="brush:php;toolbar:false;">def __call__(self, pred_cls, target_cls, pred_reg, target_reg):
    l_cls = self.cls_loss(pred_cls, target_cls)
    l_reg = self.reg_loss(pred_reg, target_reg)
    return self.w_cls * l_cls + self.w_reg * l_reg

正确用法：

loss_fn = MultiTaskLoss(w_cls=0.7, w_reg=0.3) total_loss = loss_fn(pred_cls, y_cls, pred_reg, y_reg)

注意：__call__ 接收原始预测张量和标签，不做 shape 校验——务必确保 pred_cls 维度匹配 nn.CrossEntropyLoss 要求（如 [N, C]），否则报错 Expected input batch_size to match target batch_size。

验证阶段必须分别计算各任务指标，不能只看 total_loss

训练时用加权总 loss 驱动优化，但验证时如果只打印 total_loss，根本无法判断是哪个任务退化了。曾见模型 total_loss 下降但分类准确率暴跌，只因回归 loss 主导了优化方向。

实操建议：

验证循环里独立调用各 loss 函数，记录 cls_loss.item() 和 reg_loss.item()，而非仅存 total_loss.item()
对分类任务额外计算 accuracy：acc = (pred_cls.argmax(1) == target_cls).float().mean()；对回归任务计算 MAE 或 RMSE
保存 checkpoint 的依据应是「主任务指标」（如分类 acc），而非 total_loss —— 后者可能因权重偏置持续下降，但主任务已失效

多任务最难的不是写 loss，而是让不同目标在共享参数空间里不互相拖累。权重、归一化、梯度监控、独立评估，每一步漏掉都可能让模型学成“偏科生”。