如何让PyTorch学习nn.Module外的损失函数可学习系数？

核心思路是直接创建带梯度的PyTorch张量作为可学习参数，无需将其放入nn.Module，只需将这些张量加入优化器即可参与训练。

步骤1：初始化可学习参数 #

创建requires_grad=True的张量，初始值可根据需求设置（比如默认1.0，或按任务比例初始化）：

import torch

# 初始化可学习权重，开启梯度追踪
alpha = torch.tensor(1.0, requires_grad=True)
beta = torch.tensor(1.0, requires_grad=True)
gamma = torch.tensor(1.0, requires_grad=True)
delta = torch.tensor(1.0, requires_grad=True)

# 若使用GPU，需将参数移至对应设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
alpha = alpha.to(device)
beta = beta.to(device)
gamma = gamma.to(device)
delta = delta.to(device)

步骤2：将参数加入优化器 #

把这些可学习参数传入优化器，让优化器负责更新它们：

optimizer = torch.optim.Adam([alpha, beta, gamma, delta], lr=1e-3)

步骤3：损失计算与训练流程 #

在训练循环中按公式计算总损失，执行反向传播和优化步骤：

# 假设已计算得到各分项损失loss0~loss3（均为带梯度的张量）
total_loss = alpha * loss0 + beta * loss1 + gamma * loss2 + delta * loss3

# 反向传播与参数更新
optimizer.zero_grad()
total_loss.backward()
optimizer.step()

# 可选：约束参数非负（损失权重通常为非负值）
with torch.no_grad():
    alpha.clamp_(min=1e-6)  # 避免权重为0导致梯度消失
    beta.clamp_(min=1e-6)
    gamma.clamp_(min=1e-6)
    delta.clamp_(min=1e-6)

额外约束（可选） #

如果需要权重满足和为1的约束，可在优化后添加归一化操作：

with torch.no_grad():
    sum_weights = alpha + beta + gamma + delta
    alpha.div_(sum_weights)
    beta.div_(sum_weights)
    gamma.div_(sum_weights)
    delta.div_(sum_weights)

内容的提问来源于stack exchange，提问作者user836026