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使用实例
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;接下来要执行打印语句:
执行完打印结果,程序运行结束: