从TensorFlow转PyTorch：如何构建多分支拼接的11输出模型？

错误原因分析 #

你碰到的running_mean should contain 3 elements not 21错误，核心是子模型内部的层（比如BatchNorm1d）输入维度不匹配——比如某个子模型里的BatchNorm层初始化时指定了3个特征维度，但实际传入的是21维输入，导致统计均值的维度和输入维度冲突，这说明你在定义子模型时，没有保证层之间的维度衔接正确。

正确实现方式 #

要复现TensorFlow Functional API的并行子模型结构，你需要让三个子模型各自接收完整的21维输入，分别输出指定维度后再拼接。以下是规范的PyTorch实现：

import torch
import torch.nn as nn

class ParallelModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 三个并行子模型，每个都以21维作为输入
        self.submodel1 = nn.Sequential(
            nn.Linear(21, 16),
            nn.ReLU(),
            nn.BatchNorm1d(16),
            nn.Linear(16, 4)
        )
        self.submodel2 = nn.Sequential(
            nn.Linear(21, 12),
            nn.ReLU(),
            nn.BatchNorm1d(12),
            nn.Linear(12, 3)
        )
        self.submodel3 = nn.Sequential(
            nn.Linear(21, 16),
            nn.ReLU(),
            nn.BatchNorm1d(16),
            nn.Linear(16, 4)
        )
        # 用ModuleList统一管理子模型，方便后续维护
        self.submodels = nn.ModuleList([self.submodel1, self.submodel2, self.submodel3])

    def forward(self, x):
        # x的形状为 (batch_size, 21)
        outputs = [model(x) for model in self.submodels]
        # 按特征维度拼接，得到最终的11维输出
        return torch.cat(outputs, dim=1)

# 测试模型有效性
if __name__ == "__main__":
    model = ParallelModel()
    test_input = torch.randn(8, 21)  # 批量输入：8个样本，每个21维
    output = model(test_input)
    print(f"输出形状: {output.shape}")  # 预期输出: torch.Size([8, 11])

关键注意事项 #

• 每个子模型的第一层nn.Linear必须以21作为输入维度，确保和输入数据匹配
• 子模型内部的BatchNorm层维度必须与前一层的输出维度一致（比如nn.BatchNorm1d(16)对应前一层nn.Linear(21,16)的输出维度16）
• 前向传播时要分别调用每个子模型，不能直接将所有层堆叠到一个Sequential里，否则会导致输入维度混乱

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