如何构造PyTorch自定义损失函数调整权重到目标值并解决梯度报错问题

报错原因 #

你编写的损失计算代码中，通过列表推导式提取所有权重求和结果、再调用torch.tensor()拼接为新张量的操作，会完全切断张量与原神经网络计算图的关联，新生成的张量没有梯度追踪标识，因此反向传播时无法找到梯度路径触发报错。

正确自定义损失实现 #

要实现权重向指定目标值调整的需求，直接对网络可学习参数与目标值做误差计算即可，全程保持梯度追踪链不中断即可，推荐使用MSE或者L1误差衡量权重与目标值的差异：

# 可自定义目标权重值
target_weight = 1.0
# 权重正则项系数，控制权重对齐目标值的优化力度
alpha = 0.1

for i in range(10):
    opt.zero_grad()
    out = net(features)
    # 原任务损失
    task_loss = torch.mean(torch.square(torch.tensor(5) - torch.sum(out)))
    # 自定义权重对齐损失
    weight_loss = 0.0
    for w in net.parameters():
        # 用MSE计算权重与目标值的差异，也可以换成torch.abs(w - target_weight).mean()使用L1损失
        weight_loss += torch.mean((w - target_weight) ** 2)
    # 合并两个损失做联合优化
    total_loss = task_loss + alpha * weight_loss
    print('loss:', total_loss.item())
    total_loss.backward()
    opt.step()

如果不需要保留原任务的拟合需求，只需要把权重调整到目标值，直接用weight_loss作为最终损失反向传播即可。

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