在Python中如何通过PyTorch实现多任务学习的权重自动平衡？-图灵python

Gradnorm是Pytorch中通过动态调整多任务损失权重来平衡梯度模长的轻量级方法。核心是根据每个任务更新共享参数的梯度范数，需要手动迭代更新，不能依赖optimizer.step；其实质是平衡梯度贡献而不是loss值，错误归一化loss会破坏梯度流。

GradNorm在Pytorch中实现多任务权重的自动平衡

GradNorm 它是最常用的，原理清晰，原理清晰 PyTorch 实现低成本自动权重平衡的方法不依赖于额外的模型，而只依赖于梯度范数来动态调整每个任务的损失权重。核心是使不同任务的梯度模长保持一致，避免任务主导训练。

关键步骤：记录反向传输前各任务损失的梯度范围；计算加权总损失；反向传输后，使用 GradNorm 公式更新每个任务的权重 task_weights；每次都要注意 step 后手动更新，不能靠 optimizer.step() 自动完成。

每项任务都需要单独调用 loss.backward(retain_graph=True)，或用 torch.autograd.grad() 每个任务梯度分别获取
task_weights 必须是 nn.Parameter 或 torch.tensor(..., requires_grad=False)，但是更新时要绕过 optimizer，直接赋值
建议将初始权重设置为全 1.0，学习率 alpha 通常取 1.5.太大容易震荡，太小收敛慢

为什么不能直接使用？ torch.nn.functional.normalize() 处理 loss 权重

有些人误以为是对的 loss 向量做 L2 归一化可以平衡任务，这是错误的 —— 归一化的是 loss 价值本身并不是其对参数的梯度贡献。多任务优化的目标是使每个任务与共享层参数的梯度相似，而不是让步 loss 数值接近。

典型错误现象：loss = torch.stack([loss_a, loss_b]).softmax(dim=0) @ torch.stack([loss_a, loss_b]) 这种操作看似“平衡”，实际上破坏了梯度流，导致共享层更新信号衰减，验证 loss 不降反升。

立即学习“Python免费学习笔记(深入)；

Python 3.14.2

Python 3.14.2是Python编程语言于2025年12月5日发布的稳定版本，属于3.14系列的第二次维护更新。该版本包含18个修复项目，重点解决多过程、数据和正则表达模块的回归问题，修复CVE-2025-12084等安全漏洞。这个版本标志着Python发展的一个重要里程碑，即自由线程模式(删除GIL)正式得到官方支持。

下载

softmax + 加权会掩盖真实梯度尺度的差异，特别是在任务中 loss 当权重突然增加时，强行压缩，但任务的实际训练强度下降
PyTorch 的 backward() 对 scalar loss 求导是有效的；是的 vector loss 直接 backward 会报错 grad can be implicitly created only for scalar outputs
真正起作用的是每项任务 loss 共享参数的梯度 torch.norm(torch.autograd.grad(loss_i, shared_params, retain_graph=True))

GradNorm 在共享 backbone 代码点在场景中

假设你有一个共享 CNN backbone，两个 head 分别输出 segmentation mask 和 depth map，对应 loss_seg 和 loss_depth：

# 学习权重的定义(注:不加入) optimizer）
task_weights = torch.nn.Parameter(torch.ones(2), requires_grad=False)
shared_params = list(model.backbone.parameters())
<h1>计算每个任务的梯度范数</h1><p>grad_norms = []
for loss_i in [loss_seg, loss_depth]:
grads = torch.autograd.grad(loss_i, shared_params, retain_graph=True, allow_unused=True)
grad_norm = torch.norm(torch.cat([g.flatten() for g in grads if g is not None]))
grad_norms.append(grad_norm.item())</p><h1>GradNorm 更新逻辑(每 batch 一次）</h1><p>if step % 10 == 0:  # 可隔若干 step 更新一次，避免抖动
mean_norm = sum(grad_norms) / len(grad_norms)
for i in range(2):
ratio = grad_norms[i] / mean_norm
task_weights.data[i] *= (ratio ** 1.5)  # alpha=1.5
task_weights.data /= task_weights.data.sum()  # 保持和为 1（可选）</p>

注意 allow_unused=True 防止某 head 梯度没有连接 backbone 导致崩溃；retain_graph=True 是因为多个 loss 共享相同的计算图。

比 GradNorm 较轻的替代方案：Uncertainty Weighting

如果 GradNorm 收敛不稳定，或者你只是想添加一行代码来启用权重平衡，你可以使用它 Kendall 提出的 uncertainty weighting：每项任务引入一个标量 log-variance 参数 log_var_i，把 loss 改写为 0.5 * torch.exp(-log_var_i) * loss_i + 0.5 * log_var_i。

log_var_i 是 nn.Parameter，随训练自动学习，无需手动调度
物理意义是任务预测的不确定性，不确定性高 → 任务自动减少 loss 权重
相比 GradNorm，它不依赖梯度计算，与任何梯度计算兼容 loss 结构，但对异常 loss 值更敏感(例如 depth loss 爆掉时 log_var 可能发散）
初始化建议：分类任务： log_var 初始化为回归任务初始化为0.0 1.0～2.0

真正困难的不是实现，而是判断哪个任务应该“让步”——有时候 loss 曲线看起来很平衡，但下游指标(如 mIoU 或 RMSE）相反，它变得更糟了。此时，我们必须回顾梯度分布的直方图，而不仅仅是盯着它 loss 数值。