PyTorch自定义MySequential使用add_module()时输出维度异常问题

你的问题核心出在MySequential类的forward方法里，return语句的缩进位置错误：

原代码中return X被写在了for循环的内部，这意味着循环刚执行完第一个模块（nn.LazyLinear(256)）就直接返回了结果，后面的nn.ReLU()和nn.LazyLinear(10)模块完全没有被执行，所以输出维度停在了256。

修正后的代码 #

只需要把return X移到for循环的外面，让所有模块依次执行完毕后再返回最终结果：

import torch
from torch import nn 
from torch.nn import functional as F 

class MySequential(nn.Module):
    """
    自定义顺序容器，用于按顺序添加模块
    """

    def __init__(self, *args):
        super(MySequential, self).__init__()

        for idx, module in enumerate(args):
            self.add_module(str(idx), module)

    def forward(self, X):
        """
        依次执行所有子模块
        """
        for module in self.children():
            X = module(X)
        # 将return移到循环外部
        return X


# 测试代码
input = torch.rand(500)
print('输入形状: ', input.shape)
net = MySequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))
out = net(input)
print('输出形状: ', out.shape)

修正后的运行结果 #

输入形状:  torch.Size([500])
输出形状:  torch.Size([10])

调整后所有模块会按顺序执行，最终输出维度就符合预期了。

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