PyTorch多次调用backward报错，求双输出模型分步训练解决方案

需要分步训练双输入双输出神经网络（分别优化两个输出分支，不同时训练），但运行时触发以下RuntimeError：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: 
[torch.FloatTensor [32, 2]], which is output 0 of AsStridedBackward0, is at version 2; expected version 1 instead.
Hint: the backtrace further above shows the operation that failed to compute its gradient. 
The variable in question was changed in there or anywhere later. Good luck!

错误原因 #

核心问题在于：当前代码中a1和a2来自同一次模型前向传播，第一次调用optimizer.step()会原地更新模型参数，而第二次a2.backward()依赖的计算图仍然基于更新前的参数版本，导致参数版本不匹配，触发Inplace操作错误。

解决方案 #

每次训练一个输出分支时，重新执行一次前向传播，确保每次反向传播都基于当前最新的模型参数，避免计算图依赖旧参数。同时无需设置retain_graph=True，因为两次前向传播是独立的计算图。

修改后的训练逻辑 #

def train(NN, optimizer, iteration):
    with trange(iteration) as t:
        torch.autograd.set_detect_anomaly(True)
        for i in t:
            # 训练第一个输出分支：重新前向传播，计算loss1并反向更新
            optimizer.zero_grad()
            loss1, _ = f(x1, x2, NN)
            loss1.backward()
            optimizer.step()

            # 训练第二个输出分支：再次重新前向传播，计算loss2并反向更新
            optimizer.zero_grad()
            _, loss2 = f(x1, x2, NN)
            loss2.backward()        
            optimizer.step()
            
            t.set_postfix(loss1=loss1.item(), loss2=loss2.item())

    return NN

补充说明 #

• f函数无需修改，它每次调用都会重新执行模型的前向传播，生成当前参数下的loss值。
• 每次训练分支前都要调用optimizer.zero_grad()，清空上一次的梯度，避免梯度累积。

完整可运行代码 #

from sklearn.utils.extmath import cartesian
import torch
import numpy as np
import torch.optim as optim
import torch.nn as nn
from tqdm.notebook import trange

x_1 = np.arange(0, 2 + 1e-2, 1e-5) 
x_2 = np.arange(0, 1 + 5e-2, 5e-2)
product_train = cartesian((x_1,x_2))
product_initial_condition0 = cartesian((np.array([0.]), x_2))
device = "cuda" if torch.cuda.is_available() else "cpu"
x1 = torch.tensor(product_train[:,0].astype(np.float32), requires_grad=True).to(device)
x2 = torch.tensor(product_train[:,1].astype(np.float32)).to(device)
initial_condition0 = torch.tensor(product_initial_condition0.astype(np.float32)).to(device)
initial_condition0_output = torch.tensor(product_initial_condition0[:,1].astype(np.float32)).to(device)


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(2, 32)
        self.linear2 = nn.Linear(32, 2)

    def forward(self, x1, x2):   
        inputs = torch.stack([x1,x2], dim=1)
        x = self.linear1(inputs)
        x = torch.sigmoid(x)  
        x = self.linear2(x)
        
        return x

NN = Model()
NN.to(device)

def f(x1, x2, NN):
    initial_conditon_Phi0 = NN(initial_condition0[:,0], initial_condition0[:,1])
    loss1 = ms_erorr(initial_conditon_Phi0[:,0], initial_condition0_output)
    loss2 = ms_erorr(initial_conditon_Phi0[:,1], initial_condition0_output)

    return loss1, loss2

def train(NN, optimizer, iteration):
    with trange(iteration) as t:
        torch.autograd.set_detect_anomaly(True)
        for i in t:
            # 训练第一个输出分支
            optimizer.zero_grad()
            loss1, _ = f(x1, x2, NN)
            loss1.backward()
            optimizer.step()

            # 训练第二个输出分支
            optimizer.zero_grad()
            _, loss2 = f(x1, x2, NN)
            loss2.backward()        
            optimizer.step()
            
            t.set_postfix(loss1=loss1.item(), loss2=loss2.item())

    return NN

iteration = 1000
optimizer = optim.Adam(NN.parameters(), lr=0.001)
ms_erorr = nn.MSELoss()
NN= train(NN, optimizer, iteration)

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