如何在torch.nn中禁用原地更新以支持双层优化二阶求导

方案1：使用无状态函数式调用替代原地参数更新 #

核心思路是不修改模型原始的nn.Parameter对象，全程维护独立的可微参数张量列表，每次内层迭代生成新的参数张量，前向传播时用torch.nn.utils.stateless.functional_call传入当前迭代的参数执行计算，全程不会打断计算图。
代码示例：

import torch
from torch.nn.utils import stateless

# 初始化超参数theta，需开启requires_grad
theta = torch.randn(..., requires_grad=True)
# 提取模型初始参数作为内层优化起点，转为普通张量保留计算图依赖
init_inner_params = {k: v.clone() for k, v in model.named_parameters()}

for t in range(OUT_MAX_ITR):
    # 每次外层迭代重置内层参数为初始值（可根据需求调整是否复用上一轮内层结果）
    current_inner_params = {k: v.clone() for k, v in init_inner_params.items()}
    for i in range(IN_MAX_ITR):
        # 用当前内层参数执行前向传播
        outputs = stateless.functional_call(model, current_inner_params, xtr)
        loss = compute_loss(outputs, theta)  # 损失需和超参数theta关联，保证计算图连通
        # 计算当前内层参数的梯度
        grads = torch.autograd.grad(loss, current_inner_params.values(), create_graph=True)
        # 生成更新后的新参数张量，不做原地修改
        for (k, v), g in zip(current_inner_params.items(), grads):
            current_inner_params[k] = v - inner_lr * g
    
    # 用最终收敛的内层参数计算外层损失
    out_loss = function_of_model_weights(current_inner_params, theta)
    # 反向传播得到超参数theta的梯度
    theta.grad = torch.autograd.grad(out_loss, theta)[0]
    # 更新超参数
    with torch.no_grad():
        theta -= outer_lr * theta.grad

方案2：使用支持高阶微分的优化器封装 #

如果不想手动实现参数更新逻辑，可以直接使用支持可微step的优化器实现，这类优化器的step函数不会执行原地修改，而是返回更新后的新参数张量，全程保留计算图。 #

注意事项 #

• 所有内层参数更新必须生成新的张量，不能修改原始nn.Parameter的数值或.data属性，否则会直接打断和超参数的计算图依赖
• 内层迭代步数较多时会产生较高显存占用，因为需要保存所有中间迭代的计算节点，可通过梯度检查点技术优化显存开销
• 若允许近似梯度，可以使用隐式微分方法，不需要保留完整的内层迭代计算图，能大幅降低显存占用，核心是通过求解共轭梯度得到最终权重对超参数的梯度

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