PyTorch中含模型输入梯度的损失函数参数梯度计算问题

问题原因 #

当损失函数依赖模型对输入的梯度（du/dx）时，模型中的偏置参数会因为对输入x求导后导数为0，完全不参与最终loss的计算图。此时调用grad(loss, model.parameters())时，这些偏置参数的梯度不存在，触发One of the differentiated Tensors appears to not have been used in the graph报错。

而loss1直接使用模型输出u，偏置参与了u的计算，所有参数都在计算图中，不会报错；设置bias=False后模型无偏置参数，自然也不会出现未使用的张量问题。

解决方案 #

方法1：允许未使用参数并过滤无效梯度 #

调用grad时添加allow_unused=True参数，然后过滤掉返回结果中的None（对应未参与计算的参数的梯度）：

import torch
import torch.nn as nn
from torch.autograd import grad

model = nn.Sequential(nn.Linear(1, 10), nn.Tanh(), nn.Linear(10, 1))

def loss2(x, u):
    d_u_x = grad(u, x, torch.ones_like(u), retain_graph=True, create_graph=True)[0]
    return torch.mean(d_u_x)

x = torch.randn(10, 1)
x.requires_grad_()
u = model(x)

loss = loss2(x, u)
# 添加allow_unused=True允许未使用的参数
d_loss_params = grad(loss, model.parameters(), retain_graph=True, allow_unused=True)
# 过滤掉无梯度的参数结果
valid_gradients = [grad for grad in d_loss_params if grad is not None]
# 若需要和参数对应，可同时保留参数与梯度
param_grad_pairs = [(param, grad) for param, grad in zip(model.parameters(), d_loss_params) if grad is not None]

方法2：手动筛选参与计算的可训练参数 #

针对当前场景，偏置参数未参与loss2的计算，可只提取模型的权重参数（不含偏置）计算梯度：

# 筛选所有线性层的weight参数（可训练且参与当前loss计算）
target_params = []
for module in model:
    if isinstance(module, nn.Linear):
        target_params.append(module.weight)

# 仅对筛选后的参数计算梯度
d_loss_params = grad(loss, target_params, retain_graph=True)

若模型结构复杂，也可通过named_parameters()筛选：

target_params = [param for name, param in model.named_parameters() if 'weight' in name]
d_loss_params = grad(loss, target_params, retain_graph=True)

注意事项 #

• 若模型包含其他类型可训练参数（如BatchNorm的gamma/beta），需根据实际计算逻辑判断是否参与loss的计算图，再调整筛选规则。
• allow_unused=True是通用解决方案，适合不确定哪些参数会参与计算的场景；手动筛选则更高效，适合明确知道有用参数的场景。

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