PyTorch中如何单独计算复合模型netA(netB)对应网络参数的梯度

核心实现逻辑 #

你需要的两种梯度计算场景可以通过截断计算图、临时冻结模块参数两种方式分别实现，以下是具体写法： #

场景1：仅计算损失相对于netA参数的梯度 #

你的原有写法是正确的，detach()作用于netB的输出张量，可以截断netB方向的梯度回传：

# 假设输入为x
input_x = ...
# netB输出detach后传入netA，梯度不会回传到netB
output = netA(netB(input_x).detach())
loss = loss_func(output)
loss.backward()
# 仅执行netA的优化器更新
optimizerA.step()
optimizerA.zero_grad()

场景2：仅计算损失相对于netB参数的梯度 #

需要临时冻结netA的所有参数，避免梯度更新到netA，计算完成后恢复netA的梯度权限即可：

input_x = ...
# 临时关闭netA参数的梯度记录
for param in netA.parameters():
    param.requires_grad = False

output = netA(netB(input_x))
loss = loss_func(output)
loss.backward()
# 仅执行netB的优化器更新
optimizerB.step()
optimizerB.zero_grad()

# 恢复netA的梯度记录，避免影响后续训练流程
for param in netA.parameters():
    param.requires_grad = True

原有尝试的问题说明 #

• loss_func(netA(netB).detach())的错误原因：你detach的是netA的最终输出张量，整个复合计算链路的梯度被完全截断，因此netA、netB的参数都无法获得梯度。
• netA.detach(netB)的错误原因：detach()是torch.Tensor类型的专属方法，netA是nn.Module子类的实例，不存在该方法，因此运行报错。

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