评价此页

★ ★ ★ ★ ★

beginner/basics/buildmodel_tutorial

在 Google Colab 中运行

注意

前往文末下载完整的示例代码。

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

构建神经网络#

创建日期：2021年2月9日 | 最近更新：2025年1月24日 | 最近验证：未验证

神经网络由对数据执行操作的层/模块组成。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")

Using cuda 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)

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

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

在输入上调用模型会返回一个二维张量，其中 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}")

Predicted class: tensor([5], device='cuda:0')

模型层#

让我们分解 FashionMNIST 模型中的各层。为了说明这一点，我们将取一个包含 3 张 28x28 图像的示例小批量（minibatch），看看它在通过网络时会发生什么。

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

torch.Size([3, 28, 28])

nn.Flatten#

我们初始化 nn.Flatten 层，将每个二维的 28x28 图像转换为包含 784 个像素值的连续数组（保持 dim=0 处的小批量维度）。

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

torch.Size([3, 784])

nn.Linear#

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

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

torch.Size([3, 20])

nn.ReLU#

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

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

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

Before ReLU: tensor([[ 0.4812, -0.1855, -0.5604, -0.5288,  0.1046, -0.6144, -0.0194,  0.5290,
          0.0763,  0.5170, -0.3571,  0.2859,  0.3366,  0.5851,  0.0283,  0.1084,
          0.2790, -0.1351,  0.4640, -0.4960],
        [ 0.1443, -0.3172, -0.0686, -0.2285, -0.1671, -0.0913,  0.0768,  1.0552,
          0.0984,  0.5287, -0.7454, -0.0322,  0.0910,  0.2811,  0.3954,  0.3928,
          0.3221, -0.1196,  0.1975, -0.5049],
        [ 0.3634, -0.2807, -0.3782, -0.1640,  0.1847, -0.3414,  0.1158,  1.2596,
         -0.2215,  0.6581, -0.6830,  0.0055,  0.1086,  0.6827, -0.1997,  0.1416,
          0.2429, -0.4890,  0.2843, -0.4609]], grad_fn=<AddmmBackward0>)


After ReLU: tensor([[0.4812, 0.0000, 0.0000, 0.0000, 0.1046, 0.0000, 0.0000, 0.5290, 0.0763,
         0.5170, 0.0000, 0.2859, 0.3366, 0.5851, 0.0283, 0.1084, 0.2790, 0.0000,
         0.4640, 0.0000],
        [0.1443, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0768, 1.0552, 0.0984,
         0.5287, 0.0000, 0.0000, 0.0910, 0.2811, 0.3954, 0.3928, 0.3221, 0.0000,
         0.1975, 0.0000],
        [0.3634, 0.0000, 0.0000, 0.0000, 0.1847, 0.0000, 0.1158, 1.2596, 0.0000,
         0.6581, 0.0000, 0.0055, 0.1086, 0.6827, 0.0000, 0.1416, 0.2429, 0.0000,
         0.2843, 0.0000]], grad_fn=<ReluBackward0>)

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 —— 范围在 [-infty, infty] 内的原始值 —— 这些值会被传递给 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")

Model structure: NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)


Layer: linear_relu_stack.0.weight | Size: torch.Size([512, 784]) | Values : tensor([[-0.0158,  0.0088,  0.0175,  ...,  0.0318,  0.0024,  0.0175],
        [-0.0138, -0.0299, -0.0257,  ..., -0.0250,  0.0036,  0.0299]],
       device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.0.bias | Size: torch.Size([512]) | Values : tensor([-0.0110, -0.0255], device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.2.weight | Size: torch.Size([512, 512]) | Values : tensor([[ 0.0076, -0.0259,  0.0349,  ..., -0.0224, -0.0156, -0.0226],
        [-0.0389, -0.0201,  0.0370,  ..., -0.0055, -0.0267, -0.0304]],
       device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.2.bias | Size: torch.Size([512]) | Values : tensor([0.0254, 0.0122], device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.4.weight | Size: torch.Size([10, 512]) | Values : tensor([[ 0.0088, -0.0188, -0.0149,  ..., -0.0250,  0.0217,  0.0397],
        [ 0.0042,  0.0350, -0.0065,  ..., -0.0095, -0.0337, -0.0142]],
       device='cuda:0', grad_fn=<SliceBackward0>)

Layer: linear_relu_stack.4.bias | Size: torch.Size([10]) | Values : tensor([-0.0266, -0.0178], device='cuda:0', grad_fn=<SliceBackward0>)

进一步阅读#

torch.nn API

脚本总运行时间： (0 分 0.511 秒)