PyTorch：不使用autograd计算权重梯度的方案问询

场景说明 #

模型分为两个流水线阶段：

import torch
import torch.nn as nn
import torch.nn.functional as F

class Pipe1(nn.Module):
    def __init__(self):
        super(Pipe1, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        return x

class Pipe2(nn.Module):
    def __init__(self):
        super(Pipe2, self).__init__()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

pipe = [Pipe1(), Pipe2()]

核心需求：

• 基于旧权重执行完整前向传播（Pipe1→Pipe2）并计算损失
• 在反向传播前修改Pipe1的叶子权重（如pipe[0].conv1.weight *=3）
• 基于更新后的Pipe1权重，结合旧损失对Pipe1输出的梯度计算权重梯度，无需手动处理每种层，同时优化当前两次Pipe1前向传播的冗余方案

解决方案：利用torch.autograd.functional.vjp优化流程 #

核心逻辑：先拿到损失对Pipe1输出的梯度，再用向量-雅可比乘积（VJP）直接计算该梯度对新权重的导数，避免冗余的反向传播流程。

具体代码实现：

# 假设已提前定义train_loader、optimizer_pipe1、optimizer_pipe2
for batch_idx, (data, target) in enumerate(train_loader):
    # 1. 旧权重下前向传播，计算损失及Pipe1输出的梯度
    with torch.no_grad():
        y_old = pipe[0](data)  # 旧权重的Pipe1输出，无需计算梯度
    y_old.requires_grad = True
    output = pipe[1](y_old)
    loss = F.nll_loss(output, target)
    
    # 计算损失对Pipe1输出的梯度 dy_dL
    loss.backward(retain_graph=False)
    dy_dL = y_old.grad.clone()
    y_old.grad = None  # 清理梯度缓存
    
    # 更新Pipe2参数
    optimizer_pipe2.step()
    optimizer_pipe2.zero_grad()
    
    # 2. 在无梯度上下文原地修改Pipe1权重
    with torch.no_grad():
        pipe[0].conv1.weight *= 3
    
    # 3. 用VJP计算损失对新权重的梯度
    def pipe1_forward(x):
        return pipe[0](x)
    
    _, grads = torch.autograd.functional.vjp(pipe1_forward, data, v=dy_dL)
    
    # 将梯度赋值给Pipe1参数，供优化器更新
    for param, grad in zip(pipe[0].parameters(), grads):
        param.grad = grad
    
    # 更新Pipe1参数
    optimizer_pipe1.step()
    optimizer_pipe1.zero_grad()

方案优势 #

• 无需手动推导每层梯度，PyTorch自动适配所有层类型
• 仅保留必要的两次Pipe1前向传播（旧权重一次、VJP内部新权重一次），无冗余计算
• 完全匹配需求：损失基于旧权重的前向结果，梯度基于更新后的权重计算

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