⚡ 快速入门：60分钟上手PyTorch

本教程将带你走完一个完整的机器学习流程，包括数据加载、模型构建、训练和预测。

📌 环境要求

Python 3.10+（PyTorch 2.7.0 要求）
已安装 PyTorch（参考安装指南）

📚 学习目标

读完本教程，你将学会：

加载和预处理数据
构建神经网络模型
训练模型并评估性能
保存和加载模型
使用模型进行预测

🚀 完整工作流程

数据准备 → 模型定义 → 损失函数 → 优化器 → 训练循环 → 评估 → 预测
   ↓           ↓           ↓         ↓          ↓        ↓       ↓
DataLoader  nn.Module  CrossEntropy  Adam    forward    eval   model()
                                           backward
                                           step

📦 第一步：准备数据

我们使用经典的 FashionMNIST 数据集——包含10类服装图片。

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

# 下载训练数据
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

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

# 创建数据加载器
batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

# 查看数据形状
for X, y in train_dataloader:
    print(f"X的形状: {X.shape}")  # [64, 1, 28, 28]
    print(f"y的形状: {y.shape}")  # [64]
    print(f"y的类型: {y.dtype}")  # torch.int64
    break

查看样本数据

import matplotlib.pyplot as plt

# 类别名称
labels_map = {
    0: "T恤",
    1: "裤子",
    2: "套头衫",
    3: "连衣裙",
    4: "外套",
    5: "凉鞋",
    6: "衬衫",
    7: "运动鞋",
    8: "包",
    9: "短靴",
}

# 可视化一些样本
figure = plt.figure(figsize=(10, 8))
cols, rows = 5, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.tight_layout()
plt.show()

🧠 第二步：定义模型

创建一个简单的神经网络：

# 获取设备（支持CUDA、MPS和CPU）
device = (
    "cuda" if torch.cuda.is_available()       # NVIDIA GPU
    else "mps" if torch.backends.mps.is_available()  # Apple Silicon (M1/M2/M3)
    else "cpu"                                 # CPU
)
print(f"使用设备: {device}")

# 定义模型
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()  # 将28x28图片展平为784维向量
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)  # 10个类别
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# 创建模型并移到设备
model = NeuralNetwork().to(device)
print(model)

输出：

使用设备: cuda
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)
  )
)

模型结构说明

输入图片 (1, 28, 28)
       ↓
   Flatten
       ↓
向量 (784,)
       ↓
Linear(784 → 512) + ReLU
       ↓
Linear(512 → 512) + ReLU
       ↓
Linear(512 → 10)
       ↓
输出 logits (10,)  ← 10个类别的得分

⚙️ 第三步：定义损失函数和优化器

# 损失函数：交叉熵（用于多分类）
loss_fn = nn.CrossEntropyLoss()

# 优化器：随机梯度下降
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# 或者使用更常用的Adam
# optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

💡 损失函数选择

任务类型	推荐损失函数
多分类	`nn.CrossEntropyLoss()`
二分类	`nn.BCEWithLogitsLoss()`
回归	`nn.MSELoss()`

🔄 第四步：定义训练和测试函数

def train(dataloader, model, loss_fn, optimizer):
    """训练一个epoch"""
    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:>7f}  [{current:>5d}/{size:>5d}]")


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"测试结果: \n 准确率: {(100*correct):>0.1f}%, 平均损失: {test_loss:>8f} \n")
    
    return correct

关键步骤解析

# 训练的核心5步：
pred = model(X)           # 1. 前向传播
loss = loss_fn(pred, y)   # 2. 计算损失
loss.backward()           # 3. 反向传播（计算梯度）
optimizer.step()          # 4. 更新参数
optimizer.zero_grad()     # 5. 清零梯度

⚠️ 常见错误

忘记 optimizer.zero_grad() → 梯度会累加
忘记 model.train() / model.eval() → Dropout和BatchNorm行为错误
忘记 torch.no_grad() → 浪费内存
忘记将数据移到GPU → 报错

🎯 第五步：训练模型

epochs = 5

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

print("训练完成!")

输出示例：

---------- Epoch 1 ----------
损失: 2.298389  [   64/60000]
损失: 2.290590  [ 6464/60000]
...
测试结果: 
 准确率: 34.4%, 平均损失: 2.248823 

---------- Epoch 2 ----------
...
---------- Epoch 5 ----------
...
测试结果: 
 准确率: 62.5%, 平均损失: 1.125247 

训练完成!

💾 第六步：保存模型

# 保存模型权重（推荐方式）
torch.save(model.state_dict(), "model.pth")
print("模型已保存到 model.pth")

📂 第七步：加载模型

# 创建模型实例
model = NeuralNetwork().to(device)

# 加载权重
model.load_state_dict(torch.load("model.pth", weights_only=True))
print("模型加载成功!")

🔮 第八步：使用模型预测

# 类别名称
classes = [
    "T恤",
    "裤子",
    "套头衫",
    "连衣裙",
    "外套",
    "凉鞋",
    "衬衫",
    "运动鞋",
    "包",
    "短靴",
]

# 进行预测
model.eval()

# 取一个测试样本
x, y = test_data[0][0], test_data[0][1]
print(f"真实标签: {classes[y]}")

with torch.no_grad():
    x = x.to(device)
    pred = model(x.unsqueeze(0))  # 添加batch维度
    predicted = pred.argmax(1).item()
    print(f'预测结果: {classes[predicted]}')

📊 完整代码汇总

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

# ==================== 1. 准备数据 ====================
training_data = datasets.FashionMNIST(root="data", train=True, download=True, transform=ToTensor())
test_data = datasets.FashionMNIST(root="data", train=False, download=True, transform=ToTensor())

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64)

# ==================== 2. 定义模型 ====================
device = "cuda" if torch.cuda.is_available() else "cpu"

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)
        return self.linear_relu_stack(x)

model = NeuralNetwork().to(device)

# ==================== 3. 训练配置 ====================
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# ==================== 4. 训练函数 ====================
def train(dataloader, model, loss_fn, optimizer):
    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()

def test(dataloader, model, loss_fn):
    model.eval()
    correct = 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            correct += (pred.argmax(1) == y).sum().item()
    return correct / len(dataloader.dataset)

# ==================== 5. 训练循环 ====================
for epoch in range(5):
    train(train_dataloader, model, loss_fn, optimizer)
    accuracy = test(test_dataloader, model, loss_fn)
    print(f"Epoch {epoch+1}: 准确率 {100*accuracy:.1f}%")

# ==================== 6. 保存模型 ====================
torch.save(model.state_dict(), "model.pth")
print("训练完成，模型已保存!")

🎓 学到了什么？

概念	对应代码
加载数据	`datasets.FashionMNIST()`, `DataLoader()`
定义模型	继承 `nn.Module`，实现 `forward()`
损失函数	`nn.CrossEntropyLoss()`
优化器	`torch.optim.SGD()`
训练模式	`model.train()`
评估模式	`model.eval()`, `torch.no_grad()`
前向传播	`pred = model(X)`
反向传播	`loss.backward()`
参数更新	`optimizer.step()`
梯度清零	`optimizer.zero_grad()`
保存模型	`torch.save(model.state_dict(), path)`
加载模型	`model.load_state_dict(torch.load(path))`

🚀 下一步

恭喜你完成了第一个PyTorch项目！接下来可以：

深入学习基础 → 张量操作 | 自动求导 | 数据集处理
构建更复杂的模型 → 神经网络 | CNN
尝试实战项目 → MNIST | 图像分类器

💡 官方资源推荐

本教程参考了 PyTorch 官方教程，推荐阅读以下资源获取更多内容：

🔗 PyTorch 官方入门教程 - 安装和快速开始
🔗 PyTorch 官方文档 - 完整API参考
🔗 PyTorch Tutorials - 官方教程集合
🔗 Learn the Basics - 官方基础教程
📺 Intro to PyTorch - YouTube Series - 官方视频教程