PyTorch 的 torch.nn 模块是构建神经网络的核心模块，它提供了许多用来构建和训练深度学习模型的工具。你可以通过 PyTorch 官方的 torch.nn 参考手册 来查阅更详细的 API 说明。下面是 torch.nn 中常用类和方法的概述，以及如何使用它们来构建神经网络。

1. torch.nn.Module

torch.nn.Module 是所有神经网络模块的基类。你可以通过继承 torch.nn.Module 来创建自己的神经网络模型。

主要方法：

  • forward()：定义网络的前向传播。
  • parameters()：返回模型中所有参数的生成器。
  • zero_grad()：将模型中的所有梯度归零。

示例：自定义神经网络

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)  # 输入层到隐藏层
        self.fc2 = nn.Linear(128, 10)  # 隐藏层到输出层

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # 展平输入图像
        x = torch.relu(self.fc1(x))  # 激活函数
        x = self.fc2(x)
        return x

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

2. 常用层（Layers）

torch.nn 提供了许多常用的神经网络层，你可以在模型中直接使用它们。

2.1 线性层 (Fully Connected Layer)

  • nn.Linear(in_features, out_features)：全连接层。
  • 用于连接上一层的所有神经元到当前层的每个神经元。
linear = nn.Linear(128, 64)  # 从 128 个输入神经元到 64 个输出神经元

2.2 卷积层 (Convolutional Layer)

  • nn.Conv2d(in_channels, out_channels, kernel_size)：二维卷积层。
  • 适用于图像数据。
conv = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)  # 输入 3 通道（RGB），输出 64 通道，卷积核大小为 3x3

2.3 池化层 (Pooling Layer)

  • nn.MaxPool2d(kernel_size, stride)：最大池化层。
  • 用于降低空间维度。
pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 2x2 的池化窗口，步长为 2

2.4 激活函数 (Activation Functions)

  • nn.ReLU()：ReLU 激活函数。
  • nn.Sigmoid()：Sigmoid 激活函数。
  • nn.Tanh()：Tanh 激活函数。
relu = nn.ReLU()

2.5 批归一化 (Batch Normalization)

  • nn.BatchNorm2d(num_features)：对卷积层的输出进行批量归一化。
  • num_features 是输入通道数。
batch_norm = nn.BatchNorm2d(64)  # 输入为 64 个通道的特征图

2.6 Dropout 层

  • nn.Dropout(p)：随机丢弃一部分神经元，防止过拟合。
  • p 是丢弃的概率。
dropout = nn.Dropout(p=0.5)  # 丢弃概率为 50%

3. 损失函数 (Loss Functions)

PyTorch 提供了许多常用的损失函数来训练模型。

3.1 交叉熵损失 (Cross Entropy Loss)

  • nn.CrossEntropyLoss()：用于多分类任务的损失函数。
  • 计算 softmax 和负对数似然损失。
criterion = nn.CrossEntropyLoss()

3.2 均方误差损失 (Mean Squared Error Loss)

  • nn.MSELoss()：用于回归任务的损失函数。
  • 计算输出和目标值之间的均方误差。
criterion = nn.MSELoss()

4. 优化器 (Optimizers)

虽然优化器通常在 torch.optim 中，但它们通常与 torch.nn.Module 结合使用。

4.1 随机梯度下降 (SGD)

  • optim.SGD(params, lr)：随机梯度下降优化器。
  • params 是模型参数，lr 是学习率。
optimizer = optim.SGD(model.parameters(), lr=0.01)

4.2 Adam

  • optim.Adam(params, lr)：Adam 优化器。
  • 结合了自适应学习率和动量。
optimizer = optim.Adam(model.parameters(), lr=0.001)

5. 模型训练与评估

训练和评估模型的流程通常如下：

# 设置训练模式
model.train()

# 获取输入数据和标签
inputs, labels = next(iter(train_loader))

# 清零梯度
optimizer.zero_grad()

# 前向传播
outputs = model(inputs)

# 计算损失
loss = criterion(outputs, labels)

# 反向传播
loss.backward()

# 更新参数
optimizer.step()

# 评估模型
model.eval()

6. torch.nn 参考手册链接

为了查阅 torch.nn 模块的详细文档，包括所有层、损失函数、优化器和其他工具，推荐访问 PyTorch 官方的 torch.nn 文档。

总结

  • torch.nn.Module 是所有神经网络层的基类，用户需要继承该类来构建自定义的神经网络。
  • 常用层：包括 Linear、Conv2d、ReLU、MaxPool2d 等。
  • 损失函数：包括 CrossEntropyLoss、MSELoss 等，分别适用于分类和回归任务。
  • 优化器：如 SGD、Adam 用于更新模型的参数。
  • 训练与评估：使用 model.train() 和 model.eval() 来切换训练和评估模式。

你可以在 PyTorch 官方文档 查阅完整的 API 参考，了解每个类和函数的详细用法。