如何使用PyTorch的ModuleList处理批量数据？

你遇到的问题本质是：nn.ModuleList是Python列表的封装，仅支持单个整数/切片作为索引，无法直接处理形状为(B,)的批量索引张量。当传入批量索引时，Python会尝试将整个张量转换为索引值，而只有单元素整数张量才能被转成Python整数，因此触发TypeError。

下面提供三种实用的批量处理方案：

方案1：用embedding统一管理层参数（效率最高） #

把所有Linear层的权重和偏置提取为大张量，通过embedding根据索引批量选取对应参数，手动完成线性变换：

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

N = 10 # ModuleList元素数量
H = 2  # 输入维度
B = 5  # 批量大小

class MyModel(nn.Module):
    def __init__(self, **kwargs):
        super(MyModel, self).__init__(**kwargs)
        self.list_of_nets = nn.ModuleList([nn.Linear(H, H) for _ in range(N)])
        # 将所有层的权重/偏置拼接成可批量索引的张量
        self.weight = nn.Parameter(T.stack([net.weight for net in self.list_of_nets]))
        self.bias = nn.Parameter(T.stack([net.bias for net in self.list_of_nets]))

    def forward(self, idx, x):
        # 根据索引批量选取对应层的参数
        selected_weight = F.embedding(idx, self.weight)
        selected_bias = F.embedding(idx, self.bias)
        # 批量计算线性变换
        output = T.bmm(x.unsqueeze(1), selected_weight).squeeze(1) + selected_bias
        return output

测试代码：

model = MyModel()
idx = T.randint(0, N, (B,))
x_input = T.rand((B, H))
output = model(idx, x_input)
print(output.shape) # 输出 torch.Size([5, 2])，符合预期

方案2：用vmap批量映射（PyTorch 1.12+） #

vmap可以自动将单样本逻辑映射到批量数据，无需修改参数结构：

import torch as T
import torch.nn as nn
from torch.func import vmap

N = 10 # ModuleList元素数量
H = 2  # 输入维度
B = 5  # 批量大小

class MyModel(nn.Module):
    def __init__(self, **kwargs):
        super(MyModel, self).__init__(**kwargs)
        self.list_of_nets = nn.ModuleList([nn.Linear(H, H) for _ in range(N)])

    def forward_single(self, idx, x):
        # 单样本的前向逻辑
        return self.list_of_nets[idx.item()](x)

    def forward(self, idx, x):
        # 用vmap批量处理
        return vmap(self.forward_single)(idx, x)

测试代码：

model = MyModel()
idx = T.randint(0, N, (B,))
x_input = T.rand((B, H))
output = model(idx, x_input)
print(output.shape) # 输出 torch.Size([5, 2])

方案3：循环遍历处理（简单直观） #

如果不需要追求极致效率，可直接循环处理每个样本后拼接结果：

import torch as T
import torch.nn as nn

N = 10 # ModuleList元素数量
H = 2  # 输入维度
B = 5  # 批量大小

class MyModel(nn.Module):
    def __init__(self, **kwargs):
        super(MyModel, self).__init__(**kwargs)
        self.list_of_nets = nn.ModuleList([nn.Linear(H, H) for _ in range(N)])

    def forward(self, idx, x):
        outputs = []
        for i in range(B):
            outputs.append(self.list_of_nets[idx[i]](x[i]))
        return T.stack(outputs)

• nn.ModuleList的核心作用是帮模型管理子模块的参数（自动注册到模型参数中），但索引逻辑完全遵循Python列表规则，不支持批量张量索引。
• 方案1适合大规模批量场景，效率最优；方案2代码最简洁，依赖高版本PyTorch；方案3适合小批量或调试阶段使用。

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