学习基础知识 || 快速入门 || 张量 || 数据集与数据加载器 || 转换 || 构建模型 || 自动微分 || 优化 || 保存与加载模型

构建神经网络 ¶

创建时间：2025 年 4 月 1 日 | 最后更新时间：2025 年 4 月 1 日 | 最后验证：未验证

神经网络由执行数据操作的层/模块组成。torch.nn 命名空间提供了构建您自己的神经网络所需的所有构建块。PyTorch 中的每个模块都是 nn.Module 的子类。神经网络本身也是一个模块，它由其他模块（层）组成。这种嵌套结构使得构建和管理复杂的架构变得容易。

在接下来的章节中，我们将构建一个神经网络来对 FashionMNIST 数据集中的图像进行分类。

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

获取训练设备 ¶

我们希望能够在 CUDA、MPS、MTIA 或 XPU 等加速器上训练我们的模型。如果当前加速器可用，我们将使用它。否则，我们使用 CPU。

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")

定义类 ¶

我们通过子类化 nn.Module 来定义我们的神经网络，并在 __init__ 中初始化神经网络层。每个 nn.Module 子类都在 forward 方法中实现输入数据的操作。

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

我们创建一个 NeuralNetwork 的实例，并将其移动到 device ，然后打印其结构。

model = NeuralNetwork().to(device)
print(model)

要使用该模型，我们将输入数据传递给它。这将执行模型的 forward 操作，以及一些后台操作。不要直接调用 model.forward() ！

调用模型对输入进行处理，返回一个 2 维张量，dim=0 对应每个类别的 10 个原始预测值，dim=1 对应每个输出的单个值。我们通过传递到 nn.Softmax 模块的实例来获取预测概率。

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

---

模型层 ¶

让我们分解 FashionMNIST 模型中的层。为了说明这一点，我们将取一个包含 3 张 28x28 大小的图像的样本 minibatch，并观察它通过网络时发生了什么。

input_image = torch.rand(3,28,28)
print(input_image.size())

nn.Flatten

我们初始化 nn.Flatten 层，将每个 2D 28x28 图像转换为连续的 784 个像素值数组（保持 minibatch 维度（在 dim=0 处）不变）。

flatten = nn.Flatten()
flat_image = flatten(input_image)
print(flat_image.size())

nn.Linear

线性层是一个模块，它使用其存储的权重和偏置对输入应用线性变换。

layer1 = nn.Linear(in_features=28*28, out_features=20)
hidden1 = layer1(flat_image)
print(hidden1.size())

nn.ReLU

非线性激活是创建模型输入和输出之间复杂映射的原因。它们在线性变换之后应用，以引入非线性，帮助神经网络学习各种现象。

在这个模型中，我们在线性层之间使用 nn.ReLU，但还有其他激活函数可以引入模型中的非线性。

print(f"Before ReLU: {hidden1}\n\n")
hidden1 = nn.ReLU()(hidden1)
print(f"After ReLU: {hidden1}")

nn.Sequential 序列

nn.Sequential 是一个有序的模块容器。数据将按照定义的顺序通过所有模块。您可以使用顺序容器快速构建网络，例如 seq_modules 。

seq_modules = nn.Sequential(
    flatten,
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)
input_image = torch.rand(3,28,28)
logits = seq_modules(input_image)

nn.Softmax 软最大化

神经网络的最后一层线性层返回 logits - 原始值在[-∞, ∞]之间 - 这些值被传递到 nn.Softmax 模块。logits 被缩放到[0, 1]的值，表示模型对每个类别的预测概率。 dim 参数表示值的维度，其和必须为 1。

softmax = nn.Softmax(dim=1)
pred_probab = softmax(logits)

模型参数

神经网络内部有许多层是参数化的，即具有关联的权重和偏差，这些权重和偏差在训练过程中被优化。通过子类化 nn.Module 自动跟踪模型对象内部定义的所有字段，并可以通过模型对象的 parameters() 或 named_parameters() 方法访问所有参数。

在本例中，我们遍历每个参数，并打印其大小和值预览。

print(f"Model structure: {model}\n\n")

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")

---

进一步阅读

• torch.nn API