Pytorch中torch.nn的学习

1、torch.nn介绍

torch.nn 是 PyTorch 中用于构建神经网络的核心模块，提供了构建、训练和部署神经网络所需的各种类和函数。

核心组件

(1) 层（Layers）

torch.nn 包含常见的神经网络层，例如：

线性层（全连接层）: nn.Linear(in_features, out_features)
卷积层: nn.Conv1d/2d/3d(in_channels, out_channels, kernel_size)
循环神经网络层: nn.RNN, nn.LSTM, nn.GRU
归一化层: nn.BatchNorm1d/2d/3d, nn.LayerNorm
池化层: nn.MaxPool2d, nn.AvgPool2d
激活函数: nn.ReLU, nn.Sigmoid, nn.Softmax 等

(2) 损失函数（Loss Functions）

分类任务: nn.CrossEntropyLoss, nn.BCELoss
回归任务: nn.MSELoss, nn.L1Loss
其他: nn.KLDivLoss, nn.HuberLoss 等

(3) 容器（Containers）

nn.Module: 所有神经网络模块的基类，自定义网络需继承此类。
nn.Sequential: 按顺序组合多个层（类似Keras的Sequential模型）。
nn.ModuleList: 动态存储子模块（如循环结构的层）。

关键特性

(1) 自动求导（Autograd）

torch.nn 与 PyTorch 的自动微分系统无缝集成。通过前向传播计算输出，反向传播自动计算梯度（通过loss.backward()）。

(2) 参数管理

通过parameters()方法获取模型的所有可学习参数。

支持GPU加速（.to(device)）和分布式训练。

(3) 预训练模型

通过torchvision.models或transformers库可直接加载预训练模型（如ResNet、BERT），这些模型均基于nn.Module实现。

示例代码

(1) 定义一个简单神经网络

python 复制代码

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)  # 输入784维，输出128维
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)    # 输出10维（如分类类别数）

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = Net()
print(model)

(2) 训练流程示例

python 复制代码

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 假设data和target是训练数据
output = model(data)
loss = criterion(output, target)
loss.backward()          # 反向传播
optimizer.step()         # 更新参数
optimizer.zero_grad()    # 梯度清零

扩展工具

nn.functional: 包含无状态的函数式接口（如F.relu），适合动态计算。
nn.init: 提供参数初始化方法（如nn.init.xavier_uniform_）。
torch.nn.utils: 包含剪枝（prune）、权重裁剪（clip_grad_norm_）等工具。

适用场景

计算机视觉（CNN）、自然语言处理（RNN/Transformer）、强化学习等。
支持动态计算图（适合可变长度输入或复杂控制流）。

torch.nn官方文档

二、torch.nn使用实例

2.1 简单示例

python 复制代码

import torch
from torch import nn

#继承nn.Module类，定义模型的结构
class My_module(nn.Module):
    def __init__(self):
        super().__init__()
    #前向传播，计算输出
    def forward(self, input):
        output = input * 2
        return output

#创建类实例
my_module = My_module()
#设置输入
x = torch.tensor([1.0, 2.1, 3.6, 2.5, 5.4])
#前向传播计算输出
output = my_module(x)
#打印输出
print(output)

创建继承自nn.module类的My_module类，并重新定义方法；接着创建类实例对象my_module；设置输入数据并通过前向传播计算模型输出；最后打印输出。

tip：重写方法时，如果记不清方法的名字，可以在Pycharm中找到"代码"-->"重写方法"：

运行结果：

python 复制代码

tensor([ 2.0000,  4.2000,  7.2000,  5.0000, 10.8000])

2.2 断点调试

如果要清楚程序每一步具体是如何运行的，可以设置断点，使用"调试"：

在my_module = My_module()处设置断点，并进入调试状态：

右侧是变量区和控制台。
单击调试栏中的"单步执行我的代码"

可以清楚地看到程序在我们代码中是如何运行的：

要创建实例对象，得先初始化类My_module：

执行完创建实例对象代码之后，可以在变量区看到一个实例my_moudle已经创建好了；并接着执行下一步：输入变量x

执行完之后可以看到变量区已经出现了输入变量x；接着执行下一步：计算模型输出

可以看到，为了计算模型输出，转去调用模型方法forward()，并将输入变量x传递给形参input；接着要执行方法中的计算输出，即将输入数据乘以2：

对imput乘以2之后，得到output，可以在变量区看到output了，此时程序还在方法forward中，要执行return output：

执行完return output后，程序跳出方法forward，且可以在变量区看到output和x；接下来要执行打印语句：

执行完打印结果，程序运行结束：