端到端训练两阶段模型时，分阶段用不同损失更新参数出现RuntimeError的解决方法求助

我完全懂你遇到的这个头疼问题！这个RuntimeError本质上是计算图的张量版本不匹配导致的——你第一次执行optimizer.step()更新了model1的参数，但之前计算loss2时生成的计算图，还绑定着model1参数的旧版本。等你对loss2做反向传播时，计算图里依赖的参数已经被原地修改成新版本了，PyTorch自然会报错说版本不对。

下面给你两种可行的解决思路，还有修改后的完整代码示例：

核心问题拆解 #

你的需求是分两步更新：

1. 用loss1更新model1，model2完全不动
2. 用loss2更新model2，model1完全不动

但原来的代码里，loss2是在model1更新前计算的，计算图还挂着旧的model1参数，更新后参数版本变了，反向传播就冲突了。所以我们需要把两个阶段的计算彻底分开，并且确保每个阶段只让目标模型参与梯度计算。

解决方法一：分阶段冻结模型+重新计算张量 #

这种方法适合你需要用更新后的model1输出来计算loss2的场景，步骤很清晰：

1. 第一阶段：冻结model2，只让model1的参数可更新，计算loss1并更新model1
2. 第二阶段：冻结model1，只让model2的参数可更新，重新计算features和loss2（因为model1已经更新，旧的计算图失效了），再更新model2

修改后的完整代码 #

import torch
import torch.nn as nn

# 定义第一个模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Linear(20, 10)
        self.conv2 = nn.Linear(10, 5)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x

# 定义第二个模型
class Net1(nn.Module):
    def __init__(self):
        super(Net1, self).__init__()
        self.conv1 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.conv1(x)
        return x

# 初始化模型和优化器
model1 = Net()
model2 = Net1()
optimizer_m1 = torch.optim.SGD(model1.parameters(), lr=0.1)
optimizer_m2 = torch.optim.SGD(model2.parameters(), lr=0.1)

# 注意：你的标签是连续值，CrossEntropyLoss是分类任务用的（标签是类别索引），这里改用MSELoss更合适
criterion = nn.MSELoss()

# 生成样本数据
inputs = torch.randn(2, 20)
labels = torch.randn(2, 1)

# ---------------------- 阶段1：更新model1，冻结model2 ----------------------
# 冻结model2的所有参数，不让它参与梯度计算
for param in model2.parameters():
    param.requires_grad = False
# 确保model1的参数可以被更新
for param in model1.parameters():
    param.requires_grad = True

optimizer_m1.zero_grad()
# 前向传播计算
features = model1(inputs)
outputs = model2(features)
# 计算loss1并反向传播
loss1 = criterion(outputs[0], labels[0])
loss1.backward()
# 更新model1
optimizer_m1.step()

# ---------------------- 阶段2：更新model2，冻结model1 ----------------------
# 冻结model1的所有参数
for param in model1.parameters():
    param.requires_grad = False
# 解冻model2的参数
for param in model2.parameters():
    param.requires_grad = True

optimizer_m2.zero_grad()
# 重新计算features和outputs（因为model1已经更新，用新参数生成计算图）
features = model1(inputs)
outputs = model2(features)
# 计算loss2并反向传播
loss2 = criterion(outputs, labels)
loss2.backward()
# 更新model2
optimizer_m2.step()

解决方法二：保存更新前的特征（适合用旧特征计算loss2） #

如果你希望loss2是基于model1更新前输出的特征计算的，那可以在第一阶段把特征从计算图中分离出来，第二阶段直接用这个分离后的特征，不需要重新计算model1的前向传播：

只需要修改阶段2的代码：

# ---------------------- 阶段1：更新model1，冻结model2 ----------------------
# （和方法一一样的代码，最后多一行保存分离后的特征）
features = model1(inputs)
outputs = model2(features)
loss1 = criterion(outputs[0], labels[0])
loss1.backward()
optimizer_m1.step()
# 保存分离后的特征，断开和model1计算图的关联
saved_features = features.detach()

# ---------------------- 阶段2：更新model2，冻结model1 ----------------------
for param in model1.parameters():
    param.requires_grad = False
for param in model2.parameters():
    param.requires_grad = True

optimizer_m2.zero_grad()
# 直接用保存的特征计算outputs
outputs = model2(saved_features)
loss2 = criterion(outputs, labels)
loss2.backward()
optimizer_m2.step()

额外注意点 #

1. 损失函数的选择：你的labels是用torch.randn生成的连续值，CrossEntropyLoss是给分类任务用的（标签必须是类别索引，比如0、1、2...），所以改用MSELoss（均方误差）才是正确的，不然会触发另一个错误。
2. 梯度清空：每个阶段都要单独清空当前优化器的梯度，避免之前的梯度残留影响更新效果。

备注：内容来源于stack exchange，提问作者WillWu