nn.ModuleList使用示例

最新推荐文章于 2024-07-17 10:37:33 发布

YH美洲大蠊

最新推荐文章于 2024-07-17 10:37:33 发布

阅读量355

点赞数 3

文章标签：机器学习

本文链接：https://blog.csdn.net/YHKKun/article/details/137169777

版权

本文介绍了PyTorch中nn.ModuleList的作用，它是一个存储可学习参数模块的容器，对于设计多层结构的深度学习模型提供了便利。通过实例展示如何使用nn.ModuleList构建包含多个SimpleLinear层的MLP模型，并进行前向传播操作。

摘要由CSDN通过智能技术生成

nn.ModuleList 是 PyTorch 中的一个容器类，它是 nn.Module 的子类，专门用于存储一组可学习参数的神经网络层或者其他任何继承自 nn.Module 类的模块。nn.ModuleList 为设计和实现复杂的深度学习模型提供了便利和灵活性，特别是在需要堆叠多层结构时。

import torch
import torch.nn as nn

# 定义一个简单的全连接层类
class SimpleLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super(SimpleLinear, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        return self.linear(x)

# 定义一个多层感知机模型，其中包含3个SimpleLinear层
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=3):
        super(MLP, self).__init__()
        self.layers = nn.ModuleList([SimpleLinear(input_size if i == 0 else hidden_size, hidden_size) for i in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

# 创建一个MLP模型
input_size = 10
hidden_size = 5
num_layers = 3
model = MLP(input_size, hidden_size, num_layers)

# 打印模型结构
print(model)

# 假设一个样本数据
sample_input = torch.randn(1, input_size)  # 一个批量大小为1，特征维度为10的样本

# 前向传播并输出结果
output = model(sample_input)
print("Output of MLP:", output.shape)  # 输出结果的形状

输出：

MLP(
  (layers): ModuleList(
    (0): SimpleLinear(
      (linear): Linear(in_features=10, out_features=5, bias=True)
    )
    (1-2): 2 x SimpleLinear(
      (linear): Linear(in_features=5, out_features=5, bias=True)
    )
  )
)
Output of MLP: torch.Size([1, 5])