PyTorch 学习路线:从 Tensor 到训练完整模型

学习阶段规划

阶段时间目标
Python 基础3~5 天能写循环、函数、类
NumPy 基础2 天理解矩阵运算
PyTorch Tensor3 天会张量计算和 GPU 迁移
Autograd 自动求导2 天理解梯度和 backward
nn.Module 基础3 天会定义 MLP
训练流程3 天会完整训练和评估
CNN5 天图片分类
Transformer 基础7 天理解 Attention
项目实战长期YOLO、PointNet 等

约一个月可以达到能读懂主流论文代码的水平。

为什么从线性回归开始

# 目标:给出面积和房间数,预测房价
area = [80, 100, 120, 150]
rooms = [2, 3, 3, 4]
price = [120, 180, 220, 300]

神经网络就是找一个函数 f,使得 f(area, rooms) ≈ price。PyTorch 的作用是:

  1. 定义 f 的结构(model)
  2. 自动计算 f 的梯度(autograd)
  3. 按梯度更新参数,让预测更准确(optimizer)

核心训练循环

import torch
import torch.nn as nn

# 1. 数据
x = torch.tensor([[80, 2], [100, 3], [120, 3], [150, 4]], dtype=torch.float32)
y = torch.tensor([120, 180, 220, 300], dtype=torch.float32)

# 2. 定义模型
model = nn.Sequential(
    nn.Linear(2, 16),
    nn.ReLU(),
    nn.Linear(16, 1)
)

# 3. 定义 loss 和 optimizer
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 4. 训练循环
for epoch in range(200):
    pred = model(x).squeeze()        # 前向传播
    loss = criterion(pred, y)        # 计算 loss

    optimizer.zero_grad()            # 清空上一步的梯度
    loss.backward()                  # 反向传播,计算梯度
    optimizer.step()                 # 按梯度更新参数

    if epoch % 50 == 0:
        print(f"Epoch {epoch}: loss = {loss.item():.2f}")

四步缺一不可

  • zero_grad() → 清除旧梯度(PyTorch 默认累积梯度)
  • backward() → 计算所有参数对 loss 的偏导数
  • step() → 把参数沿梯度方向更新一步

Tensor 基础操作

import torch

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.zeros(3, 4)
c = torch.randn(2, 3)       # 标准正态分布

# 形状操作
c.shape                     # torch.Size([2, 3])
c.reshape(3, 2)
c.unsqueeze(0)              # 增加维度: [1, 2, 3]
c.squeeze()                 # 去掉维度为 1 的维

# GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
c = c.to(device)

常用 Loss 函数

任务Loss调用
回归MSELossnn.MSELoss()
二分类BCELossnn.BCEWithLogitsLoss()
多分类CrossEntropynn.CrossEntropyLoss()

常用 Optimizer

# Adam:自适应学习率,大多数情况首选
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# AdamW:Adam + 权重衰减,Transformer 常用
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)

# SGD:经典,配合 momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

保存和加载模型

# 保存权重
torch.save(model.state_dict(), "model.pth")

# 加载权重
model.load_state_dict(torch.load("model.pth"))
model.eval()  # 推理时关闭 dropout 和 BatchNorm 的训练行为

入门建议用 model.state_dict() 而不是直接 torch.save(model, ...),后者会把整个模型类序列化,跨环境加载容易出问题。