如何在PyTorch中将预训练网络作为另一网络的可反向传播损失函数

cuDNN对RNN算子存在模式限制：当RNN处于eval模式时，cuDNN后端不支持反向传播计算，哪怕你已经冻结了g的所有参数、仅需要计算对输入的梯度也会触发该报错。你之前将g全局设置为eval是为了固定预训练权重、关闭dropout/BatchNorm等训练阶段才启用的逻辑，但这和RNN反向传播的要求产生了冲突。

• 保留g的train模式，仅冻结参数不切换全局运行模式，同时手动禁用g内部所有训练专属层的训练逻辑，保证g的前向输出和eval模式完全一致：

import torch
import torch.nn.functional as F

# 冻结g的所有参数，保证训练过程中g的权重不会被更新
for param in g.parameters():
    param.requires_grad = False

# 递归遍历g的所有子层，单独把dropout、BatchNorm等会受模式影响的层设为eval
def disable_training_layers(module):
    if isinstance(module, (torch.nn.Dropout, torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d)):
        module.eval()

g.apply(disable_training_layers)

• 不要将g的前向计算整体包裹在torch.no_grad()上下文内，否则会丢失计算图无法计算输入梯度。如果不需要对g(y)计算梯度，可以单独把这部分包裹在torch.no_grad()里节省显存：

# g(y)不需要计算梯度，放在no_grad里减少显存占用
with torch.no_grad():
    g_target = g(y)
# g(f(x))保留计算图，支持后续反向传播和输入梯度计算
g_output = g(f(x))
loss = F.mse_loss(g_output, g_target)

• 如果仍有cuDNN相关报错，可以临时禁用cuDNN作为兜底方案，缺点是RNN计算速度会有所下降：

torch.backends.cudnn.enabled = False

你可以直接将冻结后的g作为普通属性挂载到你的LightningModule中，Lightning不会更新requires_grad=False的参数，完全符合你将g作为固定损失函数的使用需求。需要计算dg(f(x))/dx时，直接调用torch.autograd.grad(loss, x)即可得到对应梯度。

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