Pytorch1 PyTorch 官方 QuickStart 超详细笔记｜

适合：刚入门 PyTorch、想快速跑通完整机器学习 pipeline 的同学

全文结构：数据准备 → 构建模型 → 训练优化 → 模型评估 → 保存加载 → 单张预测

一、前言

这篇笔记把 PyTorch 官方最快入门案例完整拆解，一行代码一个知识点，帮你快速掌握：

数据集怎么加载
模型怎么定义
训练循环怎么写
模型怎么保存与推理

任务：FashionMNIST 服饰图片分类（10 分类）

二、环境与依赖

python 复制代码

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

三、数据准备（Dataset + DataLoader）

1. 下载数据集

python 复制代码

# 训练集
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# 测试集
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

2. 构建 DataLoader

python 复制代码

batch_size = 64

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

四、构建模型（继承 nn.Module）

1. 选择设备

python 复制代码

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

2. 定义网络结构

python 复制代码

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)

五、训练相关配置（损失函数 + 优化器）

python 复制代码

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

六、训练函数 & 测试函数

1. 训练函数

python 复制代码

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # 前向传播
        pred = model(X)
        loss = loss_fn(pred, y)

        # 反向传播 + 更新
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

2. 测试函数

python 复制代码

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error:")
    print(f" Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

七、开始训练

python 复制代码

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

八、模型保存与加载

1. 保存模型

python 复制代码

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

2. 加载模型

python 复制代码

model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))

九、单张图片推理

python 复制代码

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
    x = x.to(device)
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

十、整体流程总结（必背）

Dataset + DataLoader 搞定数据
class 继承 nn.Module 定义模型
loss_fn + optimizer 配置训练
train 函数：前向 → loss → 反向 → 更新 → 清零
test 函数：eval() + no_grad()
save/load 完成模型持久化
eval() + no_grad() 做推理