注意
转到末尾 下载完整的示例代码。
学习基础知识 || 快速入门 || 张量 || 数据集与数据加载器 || 转换 || 构建模型 || 自动微分 || 优化 || 保存与加载模型
构建神经网络#
创建日期：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([1], 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#
这个线性层是一个模块，它使用存储的权重和偏置对输入应用线性变换。
torch.Size([3, 20])
nn.ReLU#
非线性激活会创建模型输入和输出之间的复杂映射。它们在线性变换后应用，以引入*非线性*，帮助神经网络学习各种现象。
在此模型中，我们在线性层之间使用 nn.ReLU，但还有其他激活可以为您的模型引入非线性。
Before ReLU: tensor([[ 0.1302, -0.0054, 0.0843, -0.1854, 0.2881, -0.3891, -0.0928, 0.4310,
0.0450, -0.0729, -0.7915, -0.0881, 0.6565, 0.4077, -0.1394, 0.3523,
-0.3672, 0.5399, 0.3080, 0.2837],
[ 0.5040, -0.0578, -0.0141, 0.1032, 0.5750, -0.2719, -0.1998, 0.0755,
-0.3386, -0.1990, -0.1849, -0.2550, 0.1801, 0.0531, 0.3223, 0.3852,
-0.0924, 0.2131, 0.0271, 0.3358],
[-0.0594, -0.0138, -0.1353, -0.2701, 0.1058, -0.1963, -0.3452, 0.4849,
0.1717, 0.1034, -0.2441, 0.0594, 0.4826, 0.2283, 0.1265, 0.3398,
-0.3300, 0.4241, -0.0118, 0.4783]], grad_fn=<AddmmBackward0>)
After ReLU: tensor([[0.1302, 0.0000, 0.0843, 0.0000, 0.2881, 0.0000, 0.0000, 0.4310, 0.0450,
0.0000, 0.0000, 0.0000, 0.6565, 0.4077, 0.0000, 0.3523, 0.0000, 0.5399,
0.3080, 0.2837],
[0.5040, 0.0000, 0.0000, 0.1032, 0.5750, 0.0000, 0.0000, 0.0755, 0.0000,
0.0000, 0.0000, 0.0000, 0.1801, 0.0531, 0.3223, 0.3852, 0.0000, 0.2131,
0.0271, 0.3358],
[0.0000, 0.0000, 0.0000, 0.0000, 0.1058, 0.0000, 0.0000, 0.4849, 0.1717,
0.1034, 0.0000, 0.0594, 0.4826, 0.2283, 0.1265, 0.3398, 0.0000, 0.4241,
0.0000, 0.4783]], 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.0115, 0.0128, 0.0073, ..., -0.0262, 0.0065, -0.0202],
[-0.0122, 0.0252, 0.0119, ..., -0.0075, 0.0186, -0.0262]],
device='cuda:0', grad_fn=<SliceBackward0>)
Layer: linear_relu_stack.0.bias | Size: torch.Size([512]) | Values : tensor([ 0.0080, -0.0149], device='cuda:0', grad_fn=<SliceBackward0>)
Layer: linear_relu_stack.2.weight | Size: torch.Size([512, 512]) | Values : tensor([[-0.0397, 0.0241, 0.0244, ..., 0.0294, -0.0304, 0.0229],
[ 0.0126, 0.0027, 0.0191, ..., -0.0408, 0.0015, 0.0235]],
device='cuda:0', grad_fn=<SliceBackward0>)
Layer: linear_relu_stack.2.bias | Size: torch.Size([512]) | Values : tensor([-0.0103, 0.0416], device='cuda:0', grad_fn=<SliceBackward0>)
Layer: linear_relu_stack.4.weight | Size: torch.Size([10, 512]) | Values : tensor([[-0.0362, 0.0275, 0.0178, ..., 0.0301, 0.0117, 0.0395],
[ 0.0037, 0.0395, -0.0366, ..., 0.0224, 0.0133, 0.0194]],
device='cuda:0', grad_fn=<SliceBackward0>)
Layer: linear_relu_stack.4.bias | Size: torch.Size([10]) | Values : tensor([0.0028, 0.0247], device='cuda:0', grad_fn=<SliceBackward0>)
进一步阅读#
脚本总运行时间： (0 分 0.494 秒)
