Pytorch官方教程学习笔记6：完整训练一个模型

最新推荐文章于 2025-07-04 23:12:29 发布

原创

最新推荐文章于 2025-07-04 23:12:29 发布 · 493 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #学习 #笔记

数据处理

PyTorch 有两个用于处理数据的基本工具：torch.utils.data.DataLoader 和 torch.utils.data.Dataset。Dataset 用于存储样本及其对应的标签，而 DataLoader 则为 Dataset 提供一个可迭代的包装器

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

PyTorch 提供了特定领域的库，例如 TorchText、TorchVision 和 TorchAudio，所有这些库都包含数据集。本教程中，我们将使用一个 TorchVision 数据集。

torchvision.datasets 模块包含许多真实视觉数据的 Dataset 对象，例如 CIFAR 和 COCO（完整列表请查看此处）。在本教程中，我们使用 FashionMNIST 数据集。每个 TorchVision 数据集都包含两个参数：transform 和 target_transform，分别用于修改样本和标签。

# Download training data from open datasets.
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

我们已经下载了数据集FashionMNIST训练集和测试集并分别下载入training_data和test_data。接下来我们将数据集 (Dataset) 作为参数传递给 DataLoader。DataLoader 为数据集提供一个可迭代的包装，并支持自动批处理、采样、打乱顺序以及多进程数据加载。在这里，我们定义了一个批大小为64，即 dataloader 中的每个元素将返回包含64个特征和标签的批次。

batch_size = 64

# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

输出结果：

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64

创建模型

在 PyTorch 中定义神经网络时，我们创建一个继承自 nn.Module 的类。在 __init__ 函数中定义网络的各层结构，并在 forward 函数中指定数据如何通过网络。为加速神经网络中的运算，如果有可用的 GPU 或 MPS（苹果芯片的加速器），我们会将网络移动到该设备上运行。

# Get cpu, gpu or mps device for training.
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

# Define model
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

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

运行结果：

Using cuda device
NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=51