如何在PyTorch中自动化向神经网络模型添加层与激活函数

你可以通过在自定义nn.Module的构造函数中遍历配置列表，自动构建层与激活函数序列，以下是实现代码：

import torch
from torch import nn

class NeuralNetwork(nn.Module):
    def __init__(self, connections, activations):
        super().__init__()
        layers = []
        # 遍历网络连接配置，自动生成层和激活函数
        for idx, (in_features, out_features) in enumerate(connections):
            # 添加线性层
            layers.append(nn.Linear(in_features, out_features))
            # 仅在非最后一层后添加激活函数
            if idx < len(connections) - 1:
                # 支持多激活函数循环使用，列表长度不足时自动复用
                layers.append(activations[idx % len(activations)])
        self.module_list = nn.ModuleList(layers)

    def forward(self, x):
        for module in self.module_list:
            x = module(x)
        return x

# 使用示例
connections = [(10, 30), (30, 2)]
activations = [nn.Sigmoid()]
model = NeuralNetwork(connections, activations)
print(model)

关键实现说明 #

• 参数化配置：通过connections定义每层的输入输出维度，activations定义中间层使用的激活函数实例
• 自动层构建：遍历connections逐个创建线性层，同时在每两个线性层之间插入激活函数（最后一层后不添加激活）
• 灵活的激活复用：通过idx % len(activations)实现激活函数的循环使用，比如仅传入一个激活函数时，所有中间层都会复用它
• 参数追踪：使用nn.ModuleList管理所有模块，确保PyTorch能正确识别并追踪模型参数

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