评价此页

★ ★ ★ ★ ★

beginner/basics/buildmodel_tutorial

在 Google Colab 中运行

注意

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

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

构建神经网络#

创建日期: 2021年02月09日 | 最后更新: 2025年01月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()！

将模型应用于输入会返回一个 2D 张量，其中 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([6], device='cuda:0')

模型层#

让我们分解一下 FashionMNIST 模型中的层。为了说明这一点，我们将使用一个包含 3 个 28x28 图像的样本小批量，看看当它通过网络时会发生什么。

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

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

nn.Flatten#

我们初始化 nn.Flatten 层，将每个 2D 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.0838, -0.1596,  0.2112,  0.2359, -0.3761,  0.4043, -0.4015, -0.0615,
          0.2180, -0.2139, -0.5868,  0.3046, -0.1642,  0.0262,  0.5605, -0.1890,
          0.5140, -0.2630,  0.4404,  0.1834],
        [-0.0791, -0.3621, -0.0149,  0.4168, -0.0165,  0.3271, -0.0582, -0.1739,
         -0.0459, -0.4254, -0.4844,  0.1458, -0.0997,  0.2241,  0.2173,  0.0705,
          0.2485, -0.2096,  0.1545,  0.0299],
        [ 0.0156, -0.0276,  0.1354,  0.2339, -0.4241,  0.3049, -0.4130, -0.7753,
          0.2701, -0.4634, -0.8258, -0.1060, -0.2186, -0.2815,  0.4303, -0.2589,
          0.6219, -0.0349,  0.3395,  0.0192]], grad_fn=<AddmmBackward0>)


After ReLU: tensor([[0.0838, 0.0000, 0.2112, 0.2359, 0.0000, 0.4043, 0.0000, 0.0000, 0.2180,
         0.0000, 0.0000, 0.3046, 0.0000, 0.0262, 0.5605, 0.0000, 0.5140, 0.0000,
         0.4404, 0.1834],
        [0.0000, 0.0000, 0.0000, 0.4168, 0.0000, 0.3271, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.1458, 0.0000, 0.2241, 0.2173, 0.0705, 0.2485, 0.0000,
         0.1545, 0.0299],
        [0.0156, 0.0000, 0.1354, 0.2339, 0.0000, 0.3049, 0.0000, 0.0000, 0.2701,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.4303, 0.0000, 0.6219, 0.0000,
         0.3395, 0.0192]], 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.0054,  0.0004,  0.0307,  ..., -0.0244, -0.0134,  0.0103],
        [ 0.0068, -0.0080, -0.0073,  ...,  0.0279,  0.0128,  0.0110]],
       device='cuda:0', grad_fn=<SliceBackward0>)

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

Layer: linear_relu_stack.2.weight | Size: torch.Size([512, 512]) | Values : tensor([[-0.0374,  0.0302,  0.0022,  ...,  0.0097, -0.0145,  0.0116],
        [ 0.0316,  0.0269,  0.0374,  ..., -0.0029, -0.0009,  0.0433]],
       device='cuda:0', grad_fn=<SliceBackward0>)

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

Layer: linear_relu_stack.4.weight | Size: torch.Size([10, 512]) | Values : tensor([[ 0.0218,  0.0092,  0.0375,  ..., -0.0028,  0.0192, -0.0376],
        [ 0.0414, -0.0280,  0.0212,  ..., -0.0354,  0.0215,  0.0339]],
       device='cuda:0', grad_fn=<SliceBackward0>)

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

进一步阅读#

torch.nn API

脚本总运行时间: (0 分钟 0.513 秒)