pytorch基础--torch.nn

一，PyTorch torch.nn模块（核心）

nn.Module 是 PyTorch 中所有自定义神经网络模型的基类。PyTorch 提供了许多内置的模块（类），这些模块都继承自 nn.Module，可以直接用于构建神经网络。以下是一些常用的类和它们的用途：

---

1. 基础层（Layers）

这些是神经网络的基本构建块。

类名	用途描述
nn.Linear	全连接层（线性变换），用于实现 $(y = x·W^T + b)$。
nn.Conv2d	二维卷积层，用于处理图像数据。
nn.Conv1d	一维卷积层，用于处理序列数据（如时间序列或文本）。
nn.Conv3d	三维卷积层，用于处理三维数据（如视频或医学图像）。
nn.MaxPool2d	二维最大池化层，用于下采样。
nn.AvgPool2d	二维平均池化层，用于下采样。
nn.Dropout	Dropout 层，用于防止过拟合。
nn.Dropout2d	2D Dropout 层，用于防止过拟合。
nn.BatchNorm2d	二维批归一化层，用于加速训练并提高模型稳定性。
nn.LayerNorm	层归一化，适用于小批量或动态输入。
nn.Embedding	嵌入层，用于将离散值（如单词索引）映射到连续向量空间。
nn.Flatten	展平层，将多维输入展平为一维。

---

2. 激活函数（Activation Functions）

这些类实现了常见的激活函数。

类名	用途描述
nn.ReLU	ReLU 激活函数。
nn.Sigmoid	Sigmoid 激活函数。
nn.Tanh	Tanh 激活函数。
nn.LeakyReLU	LeakyReLU 激活函数。
nn.Softmax	Softmax 激活函数，用于多分类问题的输出层。
nn.LogSoftmax	LogSoftmax 激活函数，通常与 nn.NLLLoss 结合使用。

---

3. 损失函数（Loss Functions）

这些类实现了常见的损失函数。

类名	用途描述
nn.CrossEntropyLoss	交叉熵损失，用于多分类问题。
nn.MSELoss	均方误差损失，用于回归问题。
nn.BCELoss	二分类交叉熵损失，用于二分类问题。
nn.BCEWithLogitsLoss	结合 Sigmoid 和二分类交叉熵损失，数值稳定性更好。
nn.L1Loss	L1 损失（绝对误差），用于回归问题。
nn.NLLLoss	负对数似然损失，通常与 nn.LogSoftmax 结合使用。
nn.HingeEmbeddingLoss	用于二分类或回归问题，支持间隔损失。

---

4. 容器类（Containers）

这些类用于组合多个模块。

类名	用途描述
nn.Sequential	顺序容器，按顺序执行多个模块。
nn.ModuleList	模块列表，用于动态创建子模块。
nn.ModuleDict	模块字典，用于通过名称访问子模块。

---

5. 循环神经网络（RNN）

这些类用于处理序列数据。

类名	用途描述
nn.RNN	基础循环神经网络层。
nn.LSTM	长短期记忆网络（LSTM），适用于长序列建模。
nn.GRU	门控循环单元（GRU），计算效率比 LSTM 更高。
nn.RNNCell	单个 RNN 单元，用于自定义 RNN 结构。
nn.LSTMCell	单个 LSTM 单元，用于自定义 LSTM 结构。
nn.GRUCell	单个 GRU 单元，用于自定义 GRU 结构。

---

6. Transformer 相关

这些类用于实现 Transformer 模型。

类名	用途描述
nn.Transformer	完整的 Transformer 模型。
nn.TransformerEncoder	Transformer 编码器。
nn.TransformerDecoder	Transformer 解码器。
nn.MultiheadAttention	多头注意力机制。

---

7. 工具类（Utilities）

这些类提供了额外的工具功能。

类名	用途描述
nn.Identity	恒等映射，直接返回输入。
nn.Parameter	可学习参数，通常用于自定义层。
nn.utils.clip_grad_norm_	梯度裁剪，防止梯度爆炸。
nn.utils.rnn.pack_padded_sequence	用于处理变长序列数据。
nn.utils.rnn.pad_packed_sequence	将打包的序列解包为填充后的序列。

---

8. 自定义模块模板

你可以通过继承 nn.Module 创建自定义模块。例如：

import torch.nn as nn

class CustomLayer(nn.Module):
    def __init__(self):
        super(CustomLayer, self).__init__()
        self.linear = nn.Linear(10, 5)
        self.dropout = nn.Dropout(0.5)

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

9. 例子

import torch.nn.functional as F # 激活函数等

class Net(nn.Module):  #继承PyTorch中所有神经网络模块的基类
    def __init__(self):
        super(Net, self).__init__()  # 初始化父类 nn.Module
        '''
        1. 初始化一个内部字典 _modules，用于存储子模块（如 nn.Linear、nn.Conv2d 等）。
        2. 初始化一个内部字典 _parameters，用于存储模型的可学习参数（如权重和偏置）。
        3. 初始化其他内部状态，例如钩子函数、设备信息等。
        如果你不调用 super().__init__()，这些基础设施将不会被初始化，导致模型无法正常工作。
        '''
        self.fc1 = nn.Linear(784, 128)  # 输入层 784 (28x28), 隐藏层 128
        self.fc2 = nn.Linear(128, 64)   # 隐藏层 64
        self.fc3 = nn.Linear(64, 10)    # 输出层 10 (10个数字类别)

    def forward(self, x):
        x = x.view(-1, 784)            # 展平输入图像
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)                # 最后一层不使用激活函数（配合CrossEntropyLoss）
        return F.log_softmax(x, dim=1)  # 使用log_softmax输出概率

二，建立神经网络

2.1 基本流程

import torch
import torch.nn as nn          # 神经网络模块
import torch.nn.functional as F # 激活函数等
import torch.optim as optim     # 优化器
from torchvision import datasets, transforms # 数据集和预处理

# 设置随机种子保证可重复性
torch.manual_seed(42)



# 1. 数据准备
transform = transforms.Compose([
    transforms.ToTensor(),         # 将PIL图像转为Tensor (0-1范围)
    transforms.Normalize((0.1307,), (0.3081,)) # 标准化：(input - mean)/std
])

train_dataset = datasets.MNIST(
    './data',                      # 数据存储路径
    train=True,                    # 训练集
    download=False,                # 开启或关闭下载
    transform=transform            # 应用预处理
)
test_dataset = datasets.MNIST(
    './data',
    train=False,                   # 测试集
    download=False,                # 开启或关闭下载
    transform=transform
)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=64,  # 每批加载64个样本
    shuffle=True    # 打乱数据顺序
)
test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=1000 # 测试时使用更大的批次
)



# 2. 定义神经网络模型
class Net(nn.Module):  #继承PyTorch中所有神经网络模块的基类
    def __init__(self):
        super(Net, self).__init__()  # 初始化父类 nn.Module
        '''
        1. 初始化一个内部字典 _modules，用于存储子模块（如 nn.Linear、nn.Conv2d 等）。
        2. 初始化一个内部字典 _parameters，用于存储模型的可学习参数（如权重和偏置）。
        3. 初始化其他内部状态，例如钩子函数、设备信息等。
        如果你不调用 super().__init__()，这些基础设施将不会被初始化，导致模型无法正常工作。
        '''
        self.fc1 = nn.Linear(784, 128)  # 输入层 784 (28x28), 隐藏层 128
        self.fc2 = nn.Linear(128, 64)   # 隐藏层 64
        self.fc3 = nn.Linear(64, 10)    # 输出层 10 (10个数字类别)

    def forward(self, x):
        x = x.view(-1, 784)            # 展平输入图像
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)                # 最后一层不使用激活函数（配合CrossEntropyLoss）
        return F.log_softmax(x, dim=1)  # 使用log_softmax输出概率



# 3. 初始化模型、损失函数和优化器
model = Net()                           # 初始化模型
criterion = nn.CrossEntropyLoss()       # 损失函数，自动包含softmax
optimizer = optim.Adam(model.parameters(), lr=0.001)   # 优化器



# 4. 训练过程
def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()          # 梯度清零
        output = model(data)           # 前向传播
        loss = criterion(output, target)  # 确定损失函数
        loss.backward()                # 反向传播
        optimizer.step()               # 参数更新

        if batch_idx % 100 == 0:       # 打印损失值
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}'
                  f' ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

# 5. 测试过程
def test():
    model.eval() # 切换到评估模式
    test_loss = 0
    correct = 0
    with torch.no_grad():  # 关闭梯度计算
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

# 6. 运行训练和测试
for epoch in range(1, 6):  # 训练5个epoch
    train(epoch)
    test()

# 7. 保存模型
torch.save(model.state_dict(), "mnist_model.pth")

# 加载保存的模型
model = Net()
model.load_state_dict(torch.load("mnist_model.pth"))
model.eval()

# 单个样本预测
with torch.no_grad():
    sample = test_dataset[0][0].unsqueeze(0)
    output = model(sample)
    prediction = output.argmax(dim=1).item()
    print(f"Predicted digit: {prediction}")

2.2 数据预览

import matplotlib.pyplot as plt
# 图像数据需转换为Tensor
# 获取一个样本
image, label = train_dataset[0]  # 第0个样本
# 调整张量维度顺序：PyTorch是 (C, H, W)，Matplotlib需要 (H, W, C)
image = image.permute(1, 2, 0)
# 显示图片
plt.imshow(image)
plt.title(f"Label: {label}")
plt.axis('off')
plt.show()

2.3 使用数据模型

2.3.1 加载模型权重

# 定义相同的模型架构
model = SimpleModel()

# 加载权重
model.load_state_dict(torch.load('model_weights.pth'))

# 设置模型为评估模式（关闭 Dropout 和 BatchNorm 等）
model.eval()

2.3.2 加载整个模型

# 直接加载整个模型
model = torch.load('model.pth')

# 设置模型为评估模式
model.eval()

特性	加载模型权重	加载整个模型
保存内容	仅保存模型参数（state_dict）	保存模型结构和参数
加载方式	需要先定义模型结构，再加载权重	直接加载模型
文件体积	较小	较大
代码依赖	低（只需定义相同的模型结构）	高（需要与保存时的代码完全一致）
灵活性	高（可加载到不同的模型架构中）	低（只能加载到相同的模型架构中）
适用场景	跨项目、跨版本使用，模型结构可能变化	快速实验，代码不变的情况下使用