如何在PyTorch中利用GPU计算二阶导数（解决CuDNN RNN问题）

DevOps 小满

2026-6-18

在PyTorch中基于GPU完成二阶导数计算（解决CuDNN RNN双重反向传播限制）

以下是几种可行方案，让你在GPU上完成Hessian/Fisher矩阵的二阶导数计算，无需切换到CPU：

禁用CuDNN RNN，改用PyTorch原生RNN实现
CuDNN优化的RNN为追求速度省略了部分支持双重反向传播的计算图节点。你可以在RNN前向传播时临时禁用CuDNN，强制使用PyTorch原生RNN/LSTM/GRU实现：

# 假设model包含CuDNN优化的LSTM层
with torch.backends.cudnn.flags(enabled=False):
    outputs, hidden_states = model.lstm(inputs)
# 后续损失计算、二阶导数计算正常在GPU执行
loss = compute_loss(outputs, targets)
# 执行双重反向传播计算二阶导数
first_grad = torch.autograd.grad(loss, model.parameters(), create_graph=True)
second_grad = torch.autograd.grad(sum([g.sum() for g in first_grad]), model.parameters())

该方法无需修改模型结构，仅临时关闭CuDNN对RNN的优化，代价是RNN前向速度略有下降，但二阶导数计算可正常在GPU运行。

手动实现RNN单元，完全控制计算图
如果原生RNN的性能或功能无法满足需求，可手动实现RNN/LSTM的cell逻辑。所有操作基于PyTorch基础张量运算，天然支持高阶导数，且全程在GPU运行：

class CustomLSTM(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.cells = torch.nn.ModuleList([
            torch.nn.LSTMCell(input_dim if i==0 else hidden_dim, hidden_dim).to('cuda')
            for i in range(num_layers)
        ])
    
    def forward(self, inputs):
        batch_size, seq_len, _ = inputs.shape
        h = [torch.zeros(batch_size, self.hidden_dim).to('cuda') for _ in range(self.num_layers)]
        c = [torch.zeros(batch_size, self.hidden_dim).to('cuda') for _ in range(self.num_layers)]
        outputs = []
        for t in range(seq_len):
            x_t = inputs[:, t, :]
            for i in range(self.num_layers):
                h[i], c[i] = self.cells[i](x_t, (h[i], c[i]))
                x_t = h[i]
            outputs.append(x_t)
        return torch.stack(outputs, dim=1), (h, c)

使用该自定义LSTM替代原CuDNN版本后，双重反向传播可正常在GPU执行，计算图完全可控。

利用PyTorch高阶导数API简化计算
PyTorch提供的torch.autograd.functional.hessian和torch.autograd.functional.jacobian等高阶API，可自动处理计算图的高阶导数逻辑，有时能绕过CuDNN RNN的限制。以计算损失对模型参数的Hessian为例：

def compute_loss_with_params(params, inputs, targets):
    # 将参数加载到模型
    model.load_state_dict({k: v for k, v in zip(model.state_dict().keys(), params)})
    outputs = model(inputs)
    return torch.nn.functional.mse_loss(outputs, targets)

# 提取模型参数为可求导的张量列表
param_list = [p.clone().detach().requires_grad_(True) for p in model.parameters()]
# 在GPU上计算Hessian矩阵
hessian_matrix = torch.autograd.functional.hessian(
    compute_loss_with_params,
    (param_list, inputs, targets),
    vectorize=True  # 启用向量化加速，降低内存占用
)

注意：高阶API内存消耗较高，大模型需谨慎使用，可通过vectorize=True参数优化。

拆分反向传播流程，分步计算二阶导数
将双重反向传播拆分为两步，避免直接对CuDNN RNN执行二阶反向：

先计算损失对RNN输出的一阶梯度并保留计算图；
再计算该梯度对模型参数的梯度，得到二阶导数。
示例代码：

# 第一步：计算损失对RNN输出的梯度，保留计算图
model.zero_grad()
outputs = model(inputs)
loss = compute_loss(outputs, targets)
grad_output = torch.autograd.grad(loss, outputs, create_graph=True)[0]

# 第二步：计算grad_output对模型参数的梯度（即二阶导数）
model.zero_grad()
second_order_grads = torch.autograd.grad(
    grad_output.sum(),  # 对梯度求和得到标量，方便反向传播
    model.parameters(),
    retain_graph=True
)

该方法通过拆分计算流程，绕过CuDNN RNN对双重反向的限制，同时保持GPU加速。

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