多个独立PyTorch模型参数拼接后的前向与反向传播问题

你的问题核心在于未将所有独立模型的参数纳入优化器的更新范围，同时需要确认计算图的连续性（你当前的拼接操作其实不会断裂计算图）。以下是具体解决方案：

问题分析 #

你代码中torch.cat([m().unsqueeze(0) for m in models], dim=0)这一步是没问题的：m()返回的是Model实例的nn.Parameter，属于可追踪梯度的张量，torch.cat作为可微分操作会保留梯度链路。但你缺少两个关键步骤：

1. 没有把所有Model实例的参数注册到优化器中，导致反向传播产生的梯度无法被用来更新参数；
2. 没有执行优化器的step()方法来应用梯度更新。

修改后的完整代码 #

import torch
import torch.nn as nn
import torch.optim as optim

# 独立参数模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()        
        self.params = nn.Parameter(data=torch.randn(18, 512))
        
    def forward(self):
        return self.params

# 后续示例网络
class SomeNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(18*512, 10)  # 示例层，根据你的需求调整
        
    def forward(self, x):
        # x shape: [N, 18, 512]
        x = x.flatten(1)
        return self.fc(x)

# 初始化组件
N = 10
device = 'cuda' if torch.cuda.is_available() else 'cpu'
models = [Model().to(device) for _ in range(N)]
some_network = SomeNetwork().to(device)

# 关键：收集所有独立模型的参数 + 后续网络的参数，一起加入优化器
optimizer = optim.Adam(
    [p for model in models for p in model.parameters()] + list(some_network.parameters()),
    lr=1e-3
)

# 示例损失函数
some_loss = nn.CrossEntropyLoss()

# 训练循环示例
for epoch in range(100):
    optimizer.zero_grad()  # 清零梯度
    
    # 拼接所有独立模型的参数
    params = torch.cat([m().unsqueeze(0) for m in models], dim=0)  # [10, 18, 512]
    
    # 前向传播
    y = some_network(params)
    # 示例目标标签（根据你的任务替换）
    target = torch.randint(0, 10, (N,)).to(device)
    loss = some_loss(y, target)
    
    # 反向传播 + 参数更新
    loss.backward()
    optimizer.step()
    
    # 打印训练状态
    if (epoch+1) % 10 == 0:
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

关键注意事项 #

• 参数收集：必须将所有Model实例的参数通过[p for model in models for p in model.parameters()]收集起来，加入优化器的参数组，否则这些独立参数不会被更新。
• 计算图连续性：torch.cat、unsqueeze都是PyTorch原生可微分操作，不会破坏梯度链路，拼接后的张量会完整保留原始参数的梯度信息。
• 梯度清零：每次训练迭代前必须调用optimizer.zero_grad()，避免梯度累积导致更新异常。

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