Pytorch nn.Module

一、torch.nn简介

torch.nn是 PyTorch 中用于构建神经网络的模块。它提供了一系列的类和函数，用于定义神经网络的各种层、损失函数、优化器等。

torch.nn提供的类：

• Module: 所有神经网络模型的基类，用于定义自定义神经网络模型。
• Linear: 线性层，进行线性变换。
• Conv2d: 二维卷积层。
• RNN, LSTM, GRU: 循环神经网络层，分别对应简单RNN、长短时记忆网络（LSTM）、门控循环单元（GRU）。
• BatchNorm2d: 二维批归一化层。
• CrossEntropyLoss, MSELoss: 分类交叉熵损失函数和均方误差损失函数等。
• 等等

torch.nn提供的函数：

• functional: 包含各种神经网络相关的函数，如激活函数 (relu, sigmoid, tanh 等)、池化函数 (max_pool2d, avg_pool2d)、归一化函数 (batch_norm) 等。
• init: 参数初始化函数，如常用的均匀分布初始化 (uniform_)、正态分布初始化 (normal_)、Xavier 初始化 (xavier_uniform_, xavier_normal_) 等。
• conv: 用于创建卷积层的函数。
• linear: 用于创建线性层的函数。
• dropout: 用于创建 dropout 层的函数。
• batch_norm: 用于创建批归一化层的函数。

这些类和函数提供了构建、训练和使用神经网络模型所需的基本组件和功能，使得用户可以方便地定义和管理各种类型的神经网络结构，并实现各种机器学习任务。

二、如何继承torch.nn提供的内置模块

在这个示例中，我们创建了一个名为 MyLinearLayer 的类，它继承自 nn.Linear。在 __init__ 方法中，我们通过调用 super() 来调用父类的初始化方法，并传入适当的参数。在 forward 方法中，我们可以添加自定义的前向传播逻辑，然后调用父类的 forward 方法来执行线性层的计算。

import torch
import torch.nn as nn
class MyLinearLayer(nn.Linear):
    def __init__(self, in_features, out_features):
        super(MyLinearLayer, self).__init__(in_features, out_features)
        # 这里可以添加自定义的初始化操作或其他设置
    def forward(self, x):
        # 这里可以添加自定义的前向传播逻辑
        return super(MyLinearLayer, self).forward(x)
# 创建自定义的线性层
my_linear_layer = MyLinearLayer(10, 5)

可以根据需要继承其他 PyTorch 内置模块，只需将相应的模块类作为父类即可。

三、torch.nn.Module的常用用法

1、继承和初始化：

为了定义自己的神经网络模型，通常需要创建一个子类，并将其继承自 torch.nn.Module。在子类的初始化函数 __init__() 中，可以定义模型的各个层和组件，并将它们作为模块的属性。

import torch
import torch.nn as nn
 
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(100, 50)
        self.fc2 = nn.Linear(50, 10)
 
    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

2、前向传播：

在子类中实现 forward() 方法来定义模型的前向传播过程。该方法接受输入数据 x，并根据模型的结构和参数计算输出结果。

3、参数管理：

torch.nn.Module 可以跟踪模型的所有可学习参数，并提供方便的方法来访问和管理这些参数。可以使用 parameters() 方法获取模型的所有参数，也可以使用 named_parameters() 方法获取参数及其对应的名称。

model = MyModel()
params = list(model.parameters())
for name, param in model.named_parameters():
    print(name, param.size())

4、移动模型到不同设备：

可以使用 to() 方法将模型移动到指定的设备（如 CPU 或 GPU）上进行计算。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

5、保存和加载模型：

可以使用 torch.save() 和 torch.load() 函数来保存和加载整个模型，也可以只保存和加载模型的参数。

torch.save(model.state_dict(), 'model.pth')
model.load_state_dict(torch.load('model.pth'))

6、自定义模型：

可以通过继承 torch.nn.Module 来定义自己的模型结构，可以根据任务需求自由组合各种层和组件，也可以在 forward() 方法中实现自定义的复杂计算逻辑。

等等.......

四、torch.nn会自动后向传播

torch.nn 模块中的大多数层（如线性层、卷积层等）以及损失函数都会自动支持反向传播。在使用这些层构建神经网络模型时，不需要手动实现反向传播算法。PyTorch 提供了自动求导（Autograd）机制，能够根据定义的前向传播过程自动计算梯度，并通过反向传播算法自动求解梯度。具体来说：

1. 在模型的前向传播过程中，PyTorch 会跟踪所有参与计算的张量，并构建计算图。
2. 一旦得到了输出结果，可以调用backword()方法自动计算梯度。
3. 在调用 backword()方法后，PyTorch 会沿着计算图执行反向传播算法，计算所有需要的梯度。
4. 梯度被存储在每个张量的.grad 属性中，可以用于参数更新。

   loss.backward()#自动计算损失函数关于模型参数的梯度，即反向传播过程