小土堆教你快速掌握PyTorch nn.Module：从零搭建神经网络

创作时间:

作者:

@小白创作中心

小土堆教你快速掌握PyTorch nn.Module：从零搭建神经网络

引用

CSDN

等

来源

https://blog.csdn.net/qq_27825451/article/details/90550890

https://blog.csdn.net/qq_42233059/article/details/126642856

https://blog.csdn.net/EnochChen_/article/details/144811685

https://blog.csdn.net/zhaohongfei_358/article/details/122797244

https://www.bilibili.com/video/BV1hE411t7RN/

https://www.bilibili.com/opus/763546414099726339

http://www.runoob.com/pytorch/pytorch-torch-nn-ref.html

https://www.cnblogs.com/tian777/p/15341522.html

在深度学习领域，PyTorch凭借其灵活的API和强大的动态计算图功能，成为了许多研究者和开发者首选的框架。而nn.Module作为PyTorch中定义神经网络模型的核心类，掌握它对于快速上手深度学习至关重要。今天，就让我们跟随B站知名UP主“小土堆”的教学视频，一起快速掌握nn.Module的使用方法。

什么是`nn.Module`

nn.Module是PyTorch中专门用于定义神经网络模型的模块化接口。它构建于autograd之上，可以用来定义和运行各种复杂的神经网络结构。无论是简单的线性模型，还是复杂的卷积神经网络，都可以通过继承nn.Module类来实现。

如何使用`nn.Module`

使用nn.Module定义模型主要涉及两个关键步骤：定义网络结构和实现前向传播逻辑。

1. 定义网络结构

在nn.Module的子类中，我们需要重写__init__方法来定义网络的结构。在这个方法中，你可以添加各种网络层，比如全连接层、卷积层、池化层等。这些层都是nn.Module的子类，可以直接实例化后作为属性保存。

2. 实现前向传播逻辑

前向传播逻辑通过重写forward方法来实现。在这个方法中，你需要定义输入数据如何通过你定义的网络层，最终得到输出结果。值得注意的是，你不需要显式调用forward方法，直接使用模型实例作为函数调用即可。

下面是一个简单的例子，展示了如何使用nn.Module定义一个线性模型：

import torch
from torch import nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)  # 定义一个全连接层，输入维度为10，输出维度为1

    def forward(self, x):
        return self.fc(x)  # 前向传播，直接返回全连接层的输出

# 创建模型实例
model = SimpleModel()

# 随机生成输入数据
input_data = torch.randn(1, 10)

# 调用模型进行前向传播
output = model(input_data)
print(output)

常用的子模块类

在实际应用中，我们会用到各种各样的网络层。以下是一些常用的nn.Module子类：

基础层
- nn.Linear：全连接层
- nn.Conv2d：二维卷积层
- nn.LSTM：长短期记忆网络层
激活函数
- nn.ReLU：修正线性单元
- nn.Sigmoid：Sigmoid函数
- nn.Tanh：双曲正切函数
池化层
- nn.MaxPool2d：最大池化层
- nn.AvgPool2d：平均池化层
容器类
- nn.Sequential：顺序容器
- nn.ModuleList：模块列表
- nn.ModuleDict：模块字典
Transformer相关
- nn.MultiheadAttention：多头注意力机制
- nn.TransformerEncoder：Transformer编码器
损失函数
- nn.CrossEntropyLoss：交叉熵损失
- nn.MSELoss：均方误差损失

实战案例：搭建一个简单的卷积神经网络

让我们通过一个具体的案例来加深理解。我们将使用nn.Module搭建一个简单的卷积神经网络（CNN），用于图像分类任务。

import torch
from torch import nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)  # 第一个卷积层
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # 第二个卷积层
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.fc1 = nn.Linear(64 * 8 * 8, 128)  # 全连接层
        self.fc2 = nn.Linear(128, 10)  # 输出层

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)

        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)

        x = x.view(x.size(0), -1)  # 展平处理
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 创建模型实例
model = SimpleCNN()

# 随机生成输入数据
input_data = torch.randn(1, 3, 32, 32)

# 调用模型进行前向传播
output = model(input_data)
print(output)

在这个案例中，我们定义了一个包含两个卷积层、两个池化层和两个全连接层的简单CNN模型。通过这个例子，你可以看到如何在__init__方法中定义网络层，以及如何在forward方法中实现前向传播逻辑。