PyTorch 教程
PyTorch 教程
PyTorch 是 Facebook(Meta)开源的深度学习框架,以其动态计算图、Pythonic 风格与灵活性著称,是学术界与工业界最流行的深度学习框架之一。
本教程内容
本教程从 PyTorch 的简介与安装讲起,系统覆盖:张量基础、自动求导(Autograd)、神经网络构建、数据加载、卷积/循环/Transformer 模型、优化器、训练技巧、模型保存与部署、分布式训练等全部核心主题。
PyTorch 的优势
- 动态计算图:灵活易调试,符合直觉
- Python 原生:与 NumPy 无缝衔接
- 强大生态:torchvision/torchtext/torchaudio
- 生产友好:TorchScript、ONNX、移动端支持
- 社区活跃:论文复现与研究首选
学习前提
需要 Python 基础,了解 NumPy 与基础机器学习概念(损失、梯度、优化)。有线性代数基础更佳。
学习建议
每章代码都动手运行(推荐 Google Colab 或本地 GPU)。从张量开始,理解 Autograd 后再学网络构建,循序渐进。
PyTorch 简介
PyTorch 简介
PyTorch 是一个开源深度学习框架,提供张量计算、自动求导与神经网络模块,支持 GPU 加速。
发展历史
- 2017 年:Facebook AI Research 发布,基于 Torch
- 2019 年:PyTorch 1.0 正式版,稳定 API
- 2022 年:PyTorch 2.0,引入编译模式大幅提速
- 如今:学术论文复现率最高的框架,工业界广泛使用
核心特性
# 1. 动态计算图(Dynamic Graph)
# 网络结构在运行时构建,调试像写普通 Python
# 2. 张量(Tensor)
# 类似 NumPy 数组,但支持 GPU 与自动求导
# 3. 自动求导(Autograd)
# 自动计算梯度,反向传播无需手写
# 4. 模块化(nn.Module)
# 层、损失、优化器都是模块,可组合
# 5. 生态
# torchvision(视觉)/ torchtext(文本)
# torchaudio(音频)/ TorchServe(部署)
PyTorch vs TensorFlow
PyTorch: TensorFlow:
- 动态图(默认) - 静态图为主(2.x 支持 eager)
- 调试直观 - 部署生态成熟(TF Serving)
- 学术界主流 - 工业界大规模部署
- Python 风格 - Keras 高层 API
# 选择建议:
# 研究/快速原型 → PyTorch
# 大规模生产部署 → 两者皆可(TorchServe 已成熟)
Hello World
import torch
# 创建张量
x = torch.tensor([1.0, 2.0, 3.0])
print(x)
# GPU 检查
print(torch.cuda.is_available()) # 是否有 GPU
print(torch.__version__) # 版本
典型应用场景
- 计算机视觉:图像分类、检测、分割
- 自然语言处理:LLM 训练、文本分类、翻译
- 语音:识别、合成、增强
- 强化学习、推荐系统、科学计算
PyTorch 是进入深度学习的理想框架。接下来从安装与张量基础开始,一步步构建你的第一个神经网络。
PyTorch 安装
PyTorch 安装
PyTorch 的安装根据操作系统与是否需要 GPU 而不同。推荐使用官方安装命令生成器(pytorch.org/get-started)选择配置。
CPU 版本安装
# pip 安装(CPU 版)
pip install torch torchvision torchaudio
# 验证
python -c "import torch; print(torch.__version__)"
GPU 版本安装(CUDA)
# 方式1: 官方源(默认匹配 CUDA)
pip install torch torchvision torchaudio
# 方式2: 指定 CUDA 版本(以 cu121 为例)
pip install torch torchvision torchaudio \
--index-url https://download.pytorch.org/whl/cu121
# 常见 CUDA 版本: cu118 / cu121 / cu124 等
# 根据 nvidia-smi 与驱动版本选择
验证 GPU 可用
import torch
print(torch.cuda.is_available()) # True = 可用
print(torch.cuda.device_count()) # GPU 数量
print(torch.cuda.get_device_name(0)) # GPU 型号
# 张量放到 GPU
x = torch.tensor([1.0]).cuda()
print(x.device) # cuda:0
macOS 安装(Apple Silicon)
# MPS 加速(Apple 芯片)
pip install torch torchvision torchaudio
import torch
print(torch.backends.mps.is_available()) # True
device = "mps"
# 使用
x = torch.tensor([1.0]).to(device)
conda 安装(可选)
# Anaconda/Miniconda 用户
conda install pytorch torchvision torchaudio \
-c pytorch
# 或 GPU 版
conda install pytorch torchvision torchaudio \
pytorch-cuda=12.1 -c pytorch -c nvidia
国内加速
# 使用国内镜像(如清华源)
pip install torch torchvision torchaudio \
-i https://pypi.tuna.tsinghua.edu.cn/simple
# 或阿里云镜像
# -i https://mirrors.aliyun.com/pypi/simple
常见问题
# 1. torch.cuda.is_available() 为 False
# → 检查驱动(nvidia-smi)/CUDA 版本匹配
# 2. 版本冲突
# → 用虚拟环境,避免与旧包冲突
# 3. 内存不足
# → 减小 batch size,或用 CPU 测试
# 4. 网络慢
# → 用国内镜像或离线 wheel
安装完成并验证 GPU 后,就可以开始张量基础的学习了。没有 GPU 也能学(CPU 慢一些),后续章节多数例子 CPU 可跑。
PyTorch 基础
PyTorch 基础
本节快速浏览 PyTorch 的核心 API:张量、自动求导、设备管理与基本运算,为后续章节打基础。
张量(Tensor)基础
import torch
# 创建
x = torch.tensor([1, 2, 3]) # 从列表
y = torch.zeros(2, 3) # 全零
z = torch.ones(2, 3) # 全一
r = torch.rand(2, 3) # 均匀分布
n = torch.randn(2, 3) # 正态分布
a = torch.arange(0, 10, 2) # 等差
e = torch.eye(3) # 单位阵
# 属性
x.shape # torch.Size([3])
x.dtype # torch.int64
x.device # cpu
# 转换
x.numpy() # Tensor → NumPy
torch.from_numpy(arr) # NumPy → Tensor
x.item() # 单元素 Tensor → Python 数
基本运算
a = torch.tensor([1.0, 2.0])
b = torch.tensor([3.0, 4.0])
a + b # 加法
a * b # 逐元素乘
a @ b # 点积(向量)
a.mul(b) # 等价逐元素乘
torch.sum(a) # 求和
torch.mean(a) # 均值
a.view(2, 1) # 改变形状
a.reshape(1, 2) # 重塑
# 矩阵乘法
m1 = torch.rand(2, 3)
m2 = torch.rand(3, 4)
m1 @ m2 # (2, 4)
设备与类型
device = "cuda" if torch.cuda.is_available() else "cpu"
# 移动设备
x = torch.rand(2, 2).to(device)
x = torch.rand(2, 2).cuda() # 快捷方式
# 类型转换
x.float() / x.double() / x.int()
x.to(torch.float32)
requires_grad(自动求导入口)
# 需要求导的张量设置 requires_grad=True
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = x ** 2 + 1
y.sum().backward() # 反向传播
print(x.grad) # 梯度: dy/dx = 2x → [4, 6]
上下文管理
# 推理时关闭梯度(省内存加速)
with torch.no_grad():
y = model(x)
# 训练/推理切换
model.train() # 训练模式(Dropout 等生效)
model.eval() # 推理模式
综合小例
# 线性回归的手动梯度下降(极简)
x = torch.randn(100, 1)
w_true = torch.tensor([[2.0]])
y = x @ w_true + 1 + torch.randn(100, 1) * 0.1
w = torch.randn(1, 1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
for step in range(200):
pred = x @ w + b
loss = ((pred - y) ** 2).mean()
loss.backward()
with torch.no_grad():
w -= 0.01 * w.grad
b -= 0.01 * b.grad
w.grad.zero_(); b.grad.zero_()
print(w.item(), b.item()) # 接近 2.0 和 1.0
掌握张量、运算、设备与 autograd 四大基础,就掌握了 PyTorch 的核心机制。接下来深入学习张量与神经网络。
PyTorch 张量
PyTorch 张量
张量(Tensor)是 PyTorch 的数据核心:多维数组 + GPU 加速 + 自动求导。本节系统介绍张量的创建、操作与变换。
张量创建
import torch
# 直接创建
torch.tensor([[1, 2], [3, 4]])
torch.tensor(3.14) # 标量(0维)
# 固定值
torch.zeros(3, 4) # 全零
torch.ones(3, 4) # 全一
torch.full((3, 4), 7) # 填充
torch.eye(4) # 单位阵
# 序列
torch.arange(0, 10) # [0..9]
torch.arange(0, 10, 2) # 步长2
torch.linspace(0, 1, 5) # 5个等分点
# 随机
torch.rand(3, 4) # [0,1) 均匀
torch.randn(3, 4) # 标准正态
torch.randint(0, 10, (3, 4)) # 整数随机
torch.randperm(10) # 随机排列
# 与 NumPy
import numpy as np
arr = np.array([1, 2, 3])
t = torch.from_numpy(arr) # numpy → tensor
t.numpy() # tensor → numpy
# 注意: 共享内存,修改会互相影响
形状操作
x = torch.rand(2, 3, 4)
x.shape # torch.Size([2, 3, 4])
x.size() # 等价
x.dim() # 维度数 3
x.numel() # 元素总数 24
# 重塑
x.view(6, 4) # 视图(需内存连续)
x.reshape(6, 4) # 更安全的重塑
x.reshape(-1) # 展平(24,)
x.flatten() # 展平
x.transpose(0, 1) # 交换维度
x.permute(2, 0, 1) # 任意重排维度
x.squeeze() # 去掉长度为1的维度
x.unsqueeze(0) # 增加维度(在第0位)
索引与切片
x = torch.arange(12).reshape(3, 4)
# [[0,1,2,3],[4,5,6,7],[8,9,10,11]]
x[0] # 第一行
x[0, 1] # 元素
x[:, 1] # 第一列
x[1:, :2] # 切片
x[x > 5] # 布尔索引 → [6,7,8,9,10,11]
x[torch.tensor([0, 2])] # 花式索引 → 第0、2行
数学运算
a = torch.rand(3)
b = torch.rand(3)
a + b / a - b / a * b # 逐元素
torch.add(a, b) / torch.sub / torch.mul / torch.div
torch.pow(a, 2) # 幂
torch.sqrt(a) / torch.exp / torch.log
torch.sin / torch.cos
torch.clamp(a, 0, 1) # 截断
torch.abs(a) / torch.round
# 归约
torch.sum(a) / torch.mean(a) / torch.max(a)
torch.min(a) / torch.var(a) / torch.std(a)
torch.argmax(a) / torch.argmin(a) # 索引
# 矩阵运算
A @ B # 矩阵乘
torch.matmul(A, B) # 等价
A.T # 转置
torch.inverse(A) # 逆矩阵
torch.det(A) # 行列式
拼接与拆分
a = torch.rand(2, 3)
b = torch.rand(2, 3)
torch.cat([a, b], dim=0) # 拼接 → (4, 3)
torch.cat([a, b], dim=1) # → (2, 6)
torch.stack([a, b]) # 堆叠新维度 → (2, 2, 3)
# 拆分
torch.split(torch.arange(10), 3) # 3,3,3,1
torch.chunk(torch.arange(10), 3) # 3,3,4
广播机制
# 形状不匹配时自动广播
a = torch.rand(3, 1)
b = torch.rand(1, 4)
a + b # → (3, 4),自动扩展
# 标量与张量
x + 1 # 每个元素 +1
张量操作是 PyTorch 的日常。多练习形状变换与索引切片,后续写网络与数据处理时会更顺手。
PyTorch 神经网络基础
PyTorch 神经网络基础
本节介绍用 PyTorch 构建神经网络的三大支柱:nn.Module(模块)、nn 层(线性/激活)、optim(优化器),以及标准训练循环。
nn.Module 定义网络
import torch.nn as nn
import torch.nn.functional as F
class MLP(nn.Module):
def __init__(self, in_dim, hidden, out_dim):
super().__init__()
self.fc1 = nn.Linear(in_dim, hidden)
self.fc2 = nn.Linear(hidden, out_dim)
def forward(self, x):
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
model = MLP(784, 128, 10)
print(model)
# 查看参数
for name, param in model.named_parameters():
print(name, param.shape)
常用层
nn.Linear(in, out) # 全连接
nn.Conv2d(c_in, c_out, k) # 卷积
nn.BatchNorm1d/2d # 批归一化
nn.Dropout(p) # 随机失活
nn.Embedding(n, d) # 词嵌入
nn.LSTM / nn.GRU # 循环层
nn.TransformerEncoder # Transformer
# 激活函数
nn.ReLU() / F.relu
nn.Sigmoid() / F.sigmoid
nn.Tanh()
nn.Softmax(dim=1)
nn.LeakyReLU()
损失函数
nn.MSELoss() # 均方误差(回归)
nn.CrossEntropyLoss() # 交叉熵(分类,含 softmax)
nn.BCELoss() # 二分类(需先 sigmoid)
nn.BCEWithLogitsLoss() # 二分类(内部含 sigmoid,推荐)
nn.L1Loss() # 平均绝对误差
nn.NLLLoss() # 负对数似然
优化器
import torch.optim as optim
optimizer = optim.SGD(model.parameters(), lr=0.01,
momentum=0.9)
optimizer = optim.Adam(model.parameters(), lr=0.001)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
optimizer = optim.RMSprop(model.parameters(), lr=1e-3)
标准训练循环
import torch
import torch.nn as nn
import torch.optim as optim
model = MLP(784, 128, 10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 训练循环
for epoch in range(10):
for x_batch, y_batch in train_loader:
# 1. 前向
output = model(x_batch)
loss = criterion(output, y_batch)
# 2. 清零梯度
optimizer.zero_grad()
# 3. 反向传播
loss.backward()
# 4. 更新参数
optimizer.step()
# 每个 epoch 后打印损失
print(f"epoch {epoch}: loss = {loss.item():.4f}")
训练循环四步法
# 记忆口诀:
# 1. forward: 算输出与损失
# 2. zero_grad: 清零梯度(避免累加)
# 3. backward: 反传梯度
# 4. step: 更新参数
# 顺序不可乱: zero_grad 必须在 backward 之前
设备迁移
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
# 数据也要迁移
for x_batch, y_batch in loader:
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)
...
神经网络开发 = 定义网络(nn.Module)+ 选损失 + 选优化器 + 写训练循环。这套模板适用于绝大多数任务,后续章节都是在此之上扩展。
PyTorch 第一个神经网络
PyTorch 第一个神经网络
本节用 MNIST 手写数字识别,完整实现你的第一个神经网络:数据加载 → 模型定义 → 训练 → 评估。
准备数据
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
# 数据预处理: 转张量 + 归一化
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 下载并加载 MNIST
train_set = torchvision.datasets.MNIST(
root="./data", train=True,
download=True, transform=transform)
test_set = torchvision.datasets.MNIST(
root="./data", train=False,
download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(
test_set, batch_size=256, shuffle=False)
定义模型
class Net(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28*28, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 10)
def forward(self, x):
x = x.view(x.size(0), -1) # 展平 28*28
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x) # 10 类得分
return x
model = Net()
训练函数
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
def train(epochs=5):
model.train()
for epoch in range(epochs):
total_loss = 0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg = total_loss / len(train_loader)
print(f"Epoch {epoch+1}: loss={avg:.4f}")
train()
评估函数
def evaluate():
model.eval()
correct = 0
total = 0
with torch.no_grad(): # 推理不计算梯度
for x, y in test_loader:
x, y = x.to(device), y.to(device)
out = model(x)
pred = out.argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
acc = correct / total
print(f"测试准确率: {acc*100:.2f}%")
return acc
evaluate()
# 输出示例: 测试准确率: 97.50%
预测新样本
model.eval()
with torch.no_grad():
# 取一个测试样本
x, y = test_set[0]
x = x.unsqueeze(0).to(device) # 加 batch 维度
out = model(x)
pred = out.argmax(dim=1).item()
print(f"真实值: {y}, 预测值: {pred}")
常见改进方向
# 1. 加深加宽网络(更多层/更多神经元)
# 2. 加 Dropout 防过拟合
# 3. 用 CNN 替代全连接(图像更优)
# 4. 数据增强(旋转/平移)
# 5. 学习率调度
# 6. 训练更久(epoch 更多)
恭喜!你已经完成了第一个完整的神经网络项目。这套"加载数据-定义模型-训练-评估"的流程,是之后所有深度学习项目的通用骨架。
PyTorch 数据处理与加载
PyTorch 数据处理与加载
数据是深度学习的基础。PyTorch 提供 Dataset 与 DataLoader 两大组件,实现数据的高效加载、批处理与预处理。
核心组件
Dataset # 数据集的抽象:定义如何获取单条数据
DataLoader # 加载器:批处理、打乱、多线程加载
# 流程: Dataset 提供样本 → DataLoader 批量产出
自定义 Dataset
import torch
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
def __init__(self, data, labels, transform=None):
self.data = data
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
x = self.data[idx]
y = self.labels[idx]
if self.transform:
x = self.transform(x)
return x, y
# 使用
dataset = MyDataset(features, labels)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
DataLoader 参数
DataLoader(
dataset,
batch_size=32, # 每批样本数
shuffle=True, # 打乱(训练用 True)
num_workers=4, # 并行加载进程数
drop_last=False, # 丢弃不足一批的尾部
pin_memory=True, # GPU 训练加速
collate_fn=custom_fn # 自定义批组装
)
迭代 DataLoader
for epoch in range(2):
for batch_idx, (x, y) in enumerate(loader):
print(f"epoch {epoch}, batch {batch_idx}: "
f"x={x.shape}, y={y.shape}")
# x: [batch, ...], y: [batch]
# 训练代码...
break # 只看第一批
随机采样与划分
from torch.utils.data import random_split, Subset
# 划分训练/验证集
total = len(dataset)
train_len = int(total * 0.8)
val_len = total - train_len
train_set, val_set = random_split(
dataset, [train_len, val_len])
# 按索引取子集
indices = list(range(100))
subset = Subset(dataset, indices)
# 加权采样(类别不平衡)
from torch.utils.data import WeightedRandomSampler
weights = [0.1, 0.9] * (len(dataset) // 2)
sampler = WeightedRandomSampler(weights, num_samples=5000)
loader = DataLoader(dataset, sampler=sampler)
collate_fn 自定义批处理
def collate_fn(batch):
# batch 是样本列表,自定义如何组成一个 batch
texts = [item[0] for item in batch]
labels = [item[1] for item in batch]
# 文本变长 → 填充到统一长度
max_len = max(len(t) for t in texts)
padded = torch.zeros(len(texts), max_len)
for i, t in enumerate(texts):
padded[i, :len(t)] = t
return padded, torch.tensor(labels)
loader = DataLoader(dataset, batch_size=8,
collate_fn=collate_fn)
与 torchvision 数据集的集成
import torchvision.datasets as datasets
import torchvision.transforms as transforms
transform = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])
train_set = datasets.CIFAR10(
root="./data", train=True,
transform=transform, download=True)
# 常用内置数据集:
# MNIST / FashionMNIST / CIFAR10 / CIFAR100
# ImageFolder(文件夹结构) / ImageNet(需申请)
性能建议
# 1. num_workers 设为 CPU 核数的 2-4 倍(IO 密集)
# 2. pin_memory=True 配合 GPU 训练
# 3. 预处理尽量在 transform 中做(并行)
# 4. 数据增强只在训练集用,验证集不用
# 5. 大数据集用内存映射或分块加载
Dataset/DataLoader 是数据管线的标准架构。掌握自定义 Dataset 与 DataLoader 参数,就能处理任意格式的数据。
PyTorch 线性回归
PyTorch 线性回归
线性回归是深度学习的"Hello World":预测连续值。本节用 PyTorch 实现完整的线性回归:数据生成、建模、训练、评估。
问题定义
# 目标: 学习 y = w*x + b,用数据估计 w 和 b
# 数据: 合成数据 y = 2x + 1 + 噪声
生成数据
import torch
import torch.nn as nn
import torch.optim as optim
torch.manual_seed(42)
# 真实参数
w_true, b_true = 2.0, 1.0
# 生成数据
x = torch.rand(100, 1) * 10 # [0,10)
y = w_true * x + b_true + torch.randn(100, 1) * 0.5
# 可视化数据(可选)
import matplotlib.pyplot as plt
plt.scatter(x.numpy(), y.numpy(), s=5)
plt.xlabel("x"); plt.ylabel("y")
plt.show()
定义模型
# 方式1: 用 nn.Linear
model = nn.Linear(1, 1)
# 方式2: 自定义 Module(更灵活)
class LinearModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(1, 1)
def forward(self, x):
return self.linear(x)
model = LinearModel()
# 查看初始参数
for p in model.parameters():
print(p.item()) # 随机初始化
训练
criterion = nn.MSELoss() # 均方误差
optimizer = optim.SGD(model.parameters(), lr=0.01)
losses = []
for epoch in range(500):
# 前向
pred = model(x)
loss = criterion(pred, y)
# 反向
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
if (epoch + 1) % 100 == 0:
print(f"Epoch {epoch+1}: loss={loss.item():.4f}")
# 输出示例:
# Epoch 100: loss=0.3185
# ...
# Epoch 500: loss=0.2412
评估与预测
# 学习到的参数
w_learned = model.linear.weight.item()
b_learned = model.linear.bias.item()
print(f"学习结果: w={w_learned:.3f}, b={b_learned:.3f}")
print(f"真实值: w=2.000, b=1.000")
# 预测
model.eval()
with torch.no_grad():
x_test = torch.tensor([[5.0], [7.0]])
y_pred = model(x_test)
print(y_pred)
# 应接近: 2*5+1=11, 2*7+1=15
# 损失曲线
plt.plot(losses)
plt.xlabel("epoch"); plt.ylabel("loss")
plt.yscale("log") # 对数刻度看收敛
plt.show()
多元线性回归
# 多特征: y = w1*x1 + w2*x2 + b
X = torch.rand(100, 2) * 10
w_true2 = torch.tensor([[2.0], [-1.0]])
y = X @ w_true2 + 1 + torch.randn(100, 1) * 0.5
model = nn.Linear(2, 1)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.02)
for epoch in range(500):
optimizer.zero_grad()
loss = criterion(model(X), y)
loss.backward()
optimizer.step()
print(model.weight.data) # 接近 [2, -1]
print(model.bias.item()) # 接近 1
多项式回归(非线性扩展)
# 用特征变换拟合非线性: y = x^2 + x + 1
x = torch.linspace(-3, 3, 200).reshape(-1, 1)
y = x**2 + x + 1 + torch.randn(200, 1) * 0.2
# 构造多项式特征 [x, x^2]
X_poly = torch.cat([x, x**2], dim=1)
model = nn.Linear(2, 1)
# 训练同前...
# 学到: 系数接近 [1, 1], 偏置接近 1
要点回顾
# 回归任务三要素:
# 1. 模型: nn.Linear(特征数, 输出数)
# 2. 损失: MSELoss(连续值)
# 3. 优化: SGD/Adam
# 训练循环四步: forward → zero_grad → backward → step
线性回归虽简单,却完整展示了 PyTorch 的训练机制。把这里的循环练熟,后续所有模型(CNN/RNN/Transformer)都是同一套路。
PyTorch 卷积神经网络
PyTorch 卷积神经网络
卷积神经网络(CNN)是图像处理的核心模型,通过卷积层提取局部特征,池化层降维。本节实现一个完整的 CNN 用于图像分类。
CNN 核心组件
# 卷积层: 提取特征
nn.Conv2d(in_channels, out_channels,
kernel_size, stride, padding)
# 池化层: 降维,保留主要特征
nn.MaxPool2d(kernel_size=2)
nn.AvgPool2d(kernel_size=2)
# 归一化与激活
nn.BatchNorm2d(num_features)
nn.ReLU()
# 全连接(分类头)
nn.Linear(in_features, out_features)
经典 CNN 结构
class CNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
# 特征提取部分
self.features = nn.Sequential(
# 3x32x32 输入 (CIFAR)
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(32, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2), # 32x16x16
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2), # 64x8x8
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2), # 128x4x4
)
# 分类部分
self.classifier = nn.Sequential(
nn.Flatten(), # 128*4*4=2048
nn.Linear(128 * 4 * 4, 256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, num_classes),
)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
model = CNN(num_classes=10)
计算特征图尺寸
# 卷积后尺寸公式:
# out = (in + 2*padding - kernel) // stride + 1
# 例: 32x32, kernel=3, padding=1, stride=1
# out = (32 + 2 - 3) // 1 + 1 = 32 (尺寸不变)
# 池化后: out = in // 2
# 32 → 16 → 8 → 4
训练 CNN
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
# 数据增强(提升泛化)
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
train_set = datasets.CIFAR10("./data", train=True,
transform=transform_train, download=True)
test_set = datasets.CIFAR10("./data", train=False,
transform=transform_test, download=True)
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=128, shuffle=True,
num_workers=4)
test_loader = torch.utils.data.DataLoader(
test_set, batch_size=256, num_workers=4)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(20):
model.train()
total = 0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
total += loss.item()
print(f"Epoch {epoch+1}: {total/len(train_loader):.4f}")
评估与准确率
def accuracy(loader):
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
pred = model(x).argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
print(f"测试准确率: {accuracy(test_loader)*100:.2f}%")
# CIFAR10 简单 CNN 约 80%+,加深网络可到 90%+
经典 CNN 架构
# 学习更深的架构(预训练可用):
# - LeNet: 最早的 CNN(手写数字)
# - AlexNet: 深度 CNN 开端
# - VGG: 简单堆叠 3x3 卷积
# - ResNet: 残差连接(解决退化)
# - EfficientNet: 均衡缩放
# torchvision 提供预训练模型
import torchvision.models as models
resnet = models.resnet18(pretrained=True)
CNN 设计要点
# 1. 通道数逐层增加(32→64→128),尺寸逐层减半
# 2. 卷积 padding 保持尺寸,池化减半
# 3. 最后展平接全连接分类
# 4. 加 BatchNorm 加速收敛
# 5. Dropout 防过拟合(分类头处)
CNN 是视觉任务的标配。掌握"卷积堆叠 + 池化降维 + 全连接分类"的结构,配合数据增强,就能训练出不错的图像分类器。
PyTorch 循环神经网络
PyTorch 循环神经网络
循环神经网络(RNN)处理序列数据:文本、时间序列、语音。本节实现 RNN/LSTM 用于序列建模,以情感分类为例。
RNN 原理
# RNN 逐时间步处理序列,维护隐藏状态 h:
# h_t = f(W_h * h_{t-1} + W_x * x_t + b)
# 问题: 长序列梯度消失 → LSTM/GRU 解决
# LSTM 增加"记忆单元"c,可选择遗忘/更新
LSTM 基本用法
import torch
import torch.nn as nn
# 输入: [batch, seq_len, input_size]
# 例: 10 个样本, 每样本 20 个词, 词向量 100 维
x = torch.randn(10, 20, 100)
lstm = nn.LSTM(input_size=100, hidden_size=128,
num_layers=2, batch_first=True)
out, (h_n, c_n) = lstm(x)
# out: [10, 20, 128] 每个时间步的输出
# h_n: [2, 10, 128] 最后一层各步隐藏状态
# c_n: [2, 10, 128] 记忆单元
# 取最后时间步输出(常用)
last = out[:, -1, :] # [10, 128]
GRU(更轻量)
gru = nn.GRU(input_size=100, hidden_size=128,
batch_first=True)
out, h_n = gru(x)
# GRU 无 c 单元,参数更少,速度更快
情感分类完整示例
class SentimentLSTM(nn.Module):
def __init__(self, vocab_size, embed_dim,
hidden_dim, num_classes=2):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=0)
self.lstm = nn.LSTM(embed_dim, hidden_dim,
batch_first=True,
bidirectional=True)
self.fc = nn.Linear(hidden_dim * 2, num_classes)
def forward(self, x):
# x: [batch, seq_len] (词索引)
emb = self.embedding(x) # [b, s, d]
out, _ = self.lstm(emb) # [b, s, 2h]
last = out[:, -1, :] # 最后步
return self.fc(last)
model = SentimentLSTM(
vocab_size=10000, embed_dim=128, hidden_dim=128)
训练序列数据
# 数据: 文本 → 词索引序列
# [batch, seq_len], 用 0 填充短句
x_batch = torch.tensor([
[12, 45, 3, 0, 0], # "我爱学习"
[7, 8, 9, 10, 11], # "这个电影很好看"
])
# 训练循环与之前相同
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
for x, y in train_loader:
optimizer.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.backward()
optimizer.step()
填充与掩码
# 变长序列填充后,LSTM 会对填充部分也计算
# 处理方式:
# 1. 取"真实最后时间步"(非填充)
# 2. 或用 pack_padded_sequence(高效)
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
# 按长度排序后打包
lengths = torch.tensor([5, 3, 2]) # 每句真实长度
x_packed = pack_padded_sequence(x, lengths,
batch_first=True,
enforce_sorted=False)
out_packed, _ = lstm(x_packed)
out, _ = pad_packed_sequence(out_packed, batch_first=True)
双向 RNN
# bidirectional=True 同时看前后文
# 输出维度翻倍: hidden * 2
# 适用: 分类/标注(全文可见); 生成类任务用单向
lstm = nn.LSTM(100, 128, batch_first=True,
bidirectional=True)
# out: [b, s, 256]
RNN 应用场景
# 1. 文本分类(情感/主题)
# 2. 序列标注(NER、词性)
# 3. 时间序列预测(股价/销量)
# 4. 机器翻译(Encoder-Decoder)
# 5. 语音识别
# 注意: 文本类任务现在多被 Transformer 取代
# RNN 在时间序列与轻量场景仍有价值
LSTM/GRU 是序列建模的经典方案。掌握输入输出形状([batch, seq, feature])与最后时间步取法,就能构建各类序列模型。
PyTorch 数据集
PyTorch 数据集
本节系统介绍 PyTorch 中数据集(Dataset)的体系:内置数据集、自定义数据集、以及数据集的组合与变换。
内置数据集(torchvision/torchtext)
# 视觉数据集(torchvision.datasets)
MNIST / FashionMNIST # 手写数字/服饰
CIFAR10 / CIFAR100 # 彩色图像分类
ImageFolder # 文件夹结构数据集
COCO / VOC # 检测/分割
# 文本数据集(torchtext.datasets)
IMDB / AG_NEWS # 文本分类
WikiText2 # 语言模型
# 音频数据集(torchaudio.datasets)
SpeechCommands / LibriSpeech
# 用法统一:
dataset = datasets.CIFAR10(
root="./data", train=True,
transform=transform, download=True)
ImageFolder(文件夹数据集)
# 目录结构:
# data/
# ├── train/
# │ ├── cat/ (标签0)
# │ │ ├── 1.jpg
# │ └── dog/ (标签1)
# │ ├── 1.jpg
from torchvision.datasets import ImageFolder
dataset = ImageFolder(
root="data/train",
transform=transforms.ToTensor())
# 类别映射
dataset.classes # ['cat', 'dog']
dataset.class_to_idx # {'cat': 0, 'dog': 1}
# 使用
img, label = dataset[0] # (图像, 标签)
自定义 Dataset 完整示例
import os
import torch
from torch.utils.data import Dataset
from PIL import Image
class ImageDataset(Dataset):
def __init__(self, image_dir, labels_file,
transform=None):
self.image_paths = []
self.labels = []
self.transform = transform
# 读取标签文件(假设 CSV: 文件名,标签)
with open(labels_file) as f:
for line in f:
name, label = line.strip().split(",")
self.image_paths.append(
os.path.join(image_dir, name))
self.labels.append(int(label))
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
img = Image.open(self.image_paths[idx]).convert("RGB")
label = self.labels[idx]
if self.transform:
img = self.transform(img)
return img, label
# 使用
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
dataset = ImageDataset("images/", "labels.csv",
transform=transform)
数据集变换(transform)
import torchvision.transforms as transforms
# 图像变换
transforms.Resize((224, 224))
transforms.CenterCrop(100)
transforms.RandomCrop(32, padding=4) # 随机裁剪
transforms.RandomHorizontalFlip() # 随机翻转
transforms.RandomRotation(15) # 随机旋转
transforms.ColorJitter(brightness=0.2) # 颜色扰动
transforms.ToTensor() # 转张量
transforms.Normalize(mean, std) # 归一化
# 组合
transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])
# 数据增强的意义:
# 增加样本多样性 → 提升泛化,缓解过拟合
数据集拆分
from torch.utils.data import random_split
# 按比例拆分
total = len(dataset)
train_len = int(total * 0.7)
val_len = int(total * 0.15)
test_len = total - train_len - val_len
train_set, val_set, test_set = random_split(
dataset, [train_len, val_len, test_len])
# 固定随机种子保证可复现
import torch
torch.manual_seed(42)
# random_split 的拆分顺序稳定
Dataset 进阶
# 1. 拼接多个数据集
from torch.utils.data import ConcatDataset
combined = ConcatDataset([set1, set2])
# 2. 链式数据集
from torch.utils.data import ChainDataset
# 3. 惰性加载大数据
# 不把所有数据放内存,__getitem__ 时读文件
# 4. 缓存预处理结果
# 用 lru_cache 或预处理后保存到磁盘
最佳实践
# 1. 训练/验证/测试三分离
# 2. 数据增强只用于训练集
# 3. transform 在 Dataset 中完成,保持可复用
# 4. 大文件用惰性加载
# 5. 固定种子保证可复现
# 6. 检查样本: 可视化/打印前几条数据
Dataset 是数据管线的"源头"。内置数据集开箱即用,自定义 Dataset 处理真实数据,配合 transform 完成预处理与增强。
PyTorch 数据转换
PyTorch 数据转换
数据转换(Transform)在数据进入模型前完成预处理与增强。本节系统介绍 torchvision.transforms 的用法与自定义变换。
为什么需要变换
# 1. 格式统一: 图像 → 张量 → 固定尺寸
# 2. 归一化: 加速收敛,稳定训练
# 3. 数据增强: 提升泛化
# 4. 类型转换: PIL/numpy → Tensor
# transform 在 Dataset.__getitem__ 中调用,
# 在 DataLoader 多进程加载时并行执行
常用图像变换
import torchvision.transforms as transforms
# 尺寸
transforms.Resize((224, 224))
transforms.Resize(256) # 短边缩放
transforms.CenterCrop(224)
# 随机增强
transforms.RandomCrop(32, padding=4)
transforms.RandomResizedCrop(224) # 随机裁剪缩放
transforms.RandomHorizontalFlip(p=0.5)
transforms.RandomVerticalFlip(p=0.5)
transforms.RandomRotation(degrees=15)
transforms.RandomAffine(degrees=10, translate=(0.1, 0.1))
transforms.ColorJitter(brightness=0.2,
contrast=0.2,
saturation=0.2,
hue=0.1)
transforms.RandomGrayscale(p=0.1)
# 张量化与归一化
transforms.ToTensor() # PIL/numpy → [0,1] 张量
transforms.ToPILImage() # 反向
transforms.Normalize(mean, std) # (x - mean) / std
Compose 组合
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
# 训练用增强,验证/测试不用(只归一化)
Lambda 自定义变换
# 简单变换用 Lambda
transforms.Lambda(lambda x: x / 255.0)
# 或自定义函数
def my_transform(img):
# 自定义逻辑
return img
transform = transforms.Compose([
transforms.Lambda(my_transform),
transforms.ToTensor(),
])
自定义 Transform 类
import torch
import random
class AddNoise:
"""给张量加高斯噪声"""
def __init__(self, std=0.01):
self.std = std
def __call__(self, tensor):
noise = torch.randn_like(tensor) * self.std
return tensor + noise
class RandomErase:
"""随机擦除一块区域(增强)"""
def __init__(self, p=0.5, scale=(0.02, 0.1)):
self.p = p
self.scale = scale
def __call__(self, img):
if random.random() < self.p:
# 实现随机擦除逻辑
pass
return img
transform = transforms.Compose([
transforms.ToTensor(),
AddNoise(0.005),
RandomErase(p=0.3),
])
文本与数值变换
# 文本(结合 torchtext)
from torchtext.data.utils import get_tokenizer
tokenizer = get_tokenizer("basic_english")
# 数值
import numpy as np
def normalize_numeric(x):
"""数值特征标准化(手动实现)"""
return (x - x.mean()) / (x.std() + 1e-8)
# 或直接用 torch 标准化
x = (x - x.mean(dim=0)) / x.std(dim=0)
变换最佳实践
# 1. ToTensor 后图像范围 [0,1],再 Normalize
# 2. Normalize 参数: 用数据集均值/标准差
# (可先计算,或使用数据集官方值)
# 3. 增强顺序: 几何变换在前,张量化在后
# 4. 训练/测试用不同 transform
# 5. 变换要可复现(固定种子/确定性操作)
数据变换是数据管线的"加工厂"。掌握内置变换、Compose 组合与自定义变换,就能为任何数据配置合适的预处理流程。
Pytorch torch
PyTorch torch
torch 是 PyTorch 的核心包,提供张量、运算、随机数、设备管理等功能。本节梳理 torch 包的常用接口。
模块结构
torch # 核心: 张量、运算、设备
torch.nn # 神经网络模块
torch.optim # 优化器
torch.utils # 数据加载等工具
torchvision # 视觉生态
torchtext # 文本生态
torchaudio # 音频生态
# 常用导入
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
张量创建(torch.*)
torch.tensor(...) # 从数据创建
torch.zeros/ones/eye # 特殊张量
torch.arange/linspace # 序列
torch.rand/randn/randint # 随机
torch.full # 填充
torch.empty # 未初始化
# 从 NumPy
torch.from_numpy(arr)
# 设备
torch.device("cuda:0")
torch.cuda.is_available()
数学运算
torch.add/sub/mul/div # 四则
torch.matmul / torch.bmm # 矩阵乘 / 批量矩阵乘
torch.pow/sqrt/exp/log # 幂与对数
torch.sin/cos/tan # 三角函数
torch.abs/round/floor/ceil
torch.clamp # 截断
torch.max/min/mean/sum # 归约
torch.argmax/argmin # 索引
torch.softmax(x, dim=1) # softmax
torch.sigmoid/tanh/relu # 激活
torch.norm # 范数
torch.topk(x, k) # 前 k 大
torch.sort / torch.unique
形状操作
torch.reshape / view
torch.transpose / permute
torch.squeeze / unsqueeze
torch.cat / stack / split / chunk
torch.flatten
torch.expand / repeat # 扩展
# 拼接
torch.cat([a, b], dim=0)
# 堆叠(新维度)
torch.stack([a, b], dim=0)
随机数管理
# 种子(可复现)
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
# 生成
torch.rand(3, 3)
torch.randn(3, 3) # 标准正态
torch.randint(0, 10, (3,))
torch.randperm(10) # 排列
# 随机数生成器
gen = torch.Generator().manual_seed(42)
torch.rand(3, generator=gen)
检查与调试
# 张量信息
x.shape / x.dtype / x.device
x.numel() / x.dim() / x.size()
# 打印
print(x)
print(x.tolist()) # 转 Python 列表
print(x.item()) # 单元素
# NaN 检查
torch.isnan(x).any()
torch.isinf(x).any()
# 梯度检查
x.requires_grad
x.grad
torch 常用工具
# 保存与加载
torch.save(model.state_dict(), "model.pth")
torch.load("model.pth")
# 梯度
torch.no_grad() # 上下文: 不计算梯度
torch.set_grad_enabled(False)
# 分布式与编译(进阶)
torch.compile(model) # 2.0+ 编译加速
torch.distributed # 分布式训练
# 设置线程
torch.set_num_threads(4)
torch 包是日常开发的主战场。多查文档(torch.Tensor 与 torch 函数),配合 IDE 补全,效率最高。
PyTorch torch.nn
PyTorch torch.nn
torch.nn 是神经网络模块库,提供层、容器、损失函数等组件。理解 nn 的设计,构建网络又快又稳。
nn 模块结构
nn.Module # 所有网络/层的基类
nn.Sequential # 顺序容器
nn.Linear/Conv2d # 常用层
nn.ReLU/Sigmoid # 激活(模块形式)
nn.CrossEntropyLoss # 损失函数
nn.Embedding # 词嵌入
nn.Dropout # 正则化
# 函数式接口(不保留参数)
import torch.nn.functional as F
F.relu(x) / F.softmax(x, dim=1)
nn.Module 基类
class MyNet(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(10, 20)
self.layer2 = nn.Linear(20, 5)
def forward(self, x):
x = torch.relu(self.layer1(x))
return self.layer2(x)
model = MyNet()
# Module 自动管理:
model.parameters() # 所有参数迭代器
model.named_parameters() # (名字, 参数)
model.state_dict() # 参数字典
model.to("cuda") # 移动设备
model.train() / model.eval() # 模式切换
model.zero_grad() # 清零梯度
model.children() # 子模块
model.apply(fn) # 递归应用函数
nn.Sequential
# 顺序网络
model = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(256, 10),
)
# 命名版本
model = nn.Sequential(
fc1=nn.Linear(784, 256),
act=nn.ReLU(),
fc2=nn.Linear(256, 10),
)
model.fc1.weight.shape
常用层速查
# 全连接
nn.Linear(in, out, bias=True)
# 卷积
nn.Conv1d/2d/3d(in_c, out_c, kernel, stride, padding)
nn.ConvTranspose2d(...) # 转置卷积(上采样)
# 池化
nn.MaxPool2d(k) / nn.AvgPool2d(k)
nn.AdaptiveAvgPool2d((1, 1)) # 自适应池化到 1x1
# 归一化
nn.BatchNorm1d/2d(num_features)
nn.LayerNorm(normalized_shape)
nn.GroupNorm(num_groups, num_channels)
# 循环
nn.RNN / nn.LSTM / nn.GRU(input, hidden)
# 其他
nn.Embedding(num_emb, dim)
nn.Dropout(p)
nn.Flatten()
nn.Softmax(dim=1) / nn.LogSoftmax(dim=1)
激活函数
# 模块形式(用于 Sequential)
nn.ReLU() / nn.LeakyReLU(0.1)
nn.Sigmoid() / nn.Tanh()
nn.Softmax(dim=1)
nn.GELU() # Transformer 常用
nn.SiLU() / nn.Swish
# 函数形式(代码中直接用)
F.relu(x) / F.softmax(x, dim=1) / F.sigmoid(x)
损失函数
# 分类
nn.CrossEntropyLoss() # 多分类(推荐)
nn.BCEWithLogitsLoss() # 二分类(推荐)
nn.NLLLoss() # 与 LogSoftmax 搭配
# 回归
nn.MSELoss() / nn.L1Loss()
nn.SmoothL1Loss() # Huber,抗离群点
# 其他
nn.KLDivLoss() # KL 散度(蒸馏)
nn.CosineSimilarity() # 余弦相似
# 用法
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)
自定义层/损失
# 自定义层(继承 Module)
class MyLayer(nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
self.weight = nn.Parameter(
torch.randn(in_dim, out_dim))
self.bias = nn.Parameter(torch.zeros(out_dim))
def forward(self, x):
return x @ self.weight + self.bias
# 自定义损失(普通函数即可)
def my_loss(pred, target):
return torch.abs(pred - target).mean()
nn 设计要点
# 1. 有参数的操作用 nn 模块,无参数用 F 函数
# 2. 子模块赋值给 self.xxx 即被自动管理
# 3. forward 决定前向计算
# 4. state_dict 保存/加载参数
# 5. 自定义层继承 nn.Module 并注册 nn.Parameter
torch.nn 是构建网络的"积木库"。掌握 Module 生命周期与常用层,任何网络都能组装出来。
Transformer 模型
Transformer 模型
Transformer 是现代 NLP 的核心架构(Attention Is All You Need, 2017),也是 GPT、BERT 等大模型的基础。本节介绍其原理与 PyTorch 实现。
核心思想
# 抛弃循环结构,用"注意力机制"处理序列
# 优点: 并行计算(快)、长距离依赖(强)
# 三大组件:
# 1. Self-Attention: 序列内相互关注
# 2. Multi-Head Attention: 多视角注意力
# 3. Position Encoding: 注入位置信息
注意力机制
# 注意力公式:
# Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V
# Q(查询)/K(键)/V(值)由输入线性变换得到
# QK^T: 计算各位置相关性
# /sqrt(d_k): 缩放防梯度消失
# softmax: 归一化成权重
# × V: 加权求和
# 直觉: 每个词关注序列中相关的词,
# 聚合它们的信息作为自己的新表示
PyTorch 手动实现注意力
import torch
import torch.nn as nn
import torch.nn.functional as F
def scaled_dot_product_attention(q, k, v, mask=None):
# q/k/v: [batch, heads, seq, dim]
d_k = q.size(-1)
scores = q @ k.transpose(-2, -1) / (d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
weights = F.softmax(scores, dim=-1)
return weights @ v, weights
多头注意力实现
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.w_out = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
batch, seq, _ = x.shape
# 线性变换并分头
q = self.w_q(x).view(batch, seq,
self.n_heads, self.d_k)
q = q.transpose(1, 2) # [b, h, s, d]
k = self.w_k(x).view(batch, seq,
self.n_heads, self.d_k).transpose(1, 2)
v = self.w_v(x).view(batch, seq,
self.n_heads, self.d_k).transpose(1, 2)
attn, _ = scaled_dot_product_attention(q, k, v, mask)
# 合并头
attn = attn.transpose(1, 2).contiguous()
attn = attn.view(batch, seq, self.d_model)
return self.dropout(self.w_out(attn))
位置编码
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2)
* (-torch.log(torch.tensor(10000.0))
/ d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer("pe", pe.unsqueeze(0))
def forward(self, x):
# x: [batch, seq, d_model]
return x + self.pe[:, :x.size(1)]
Transformer 编码器块
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads,
d_ff, dropout=0.1):
super().__init__()
self.attn = MultiHeadAttention(d_model, n_heads)
self.norm1 = nn.LayerNorm(d_model)
# 前馈网络
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.ReLU(),
nn.Linear(d_ff, d_model),
)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
# 残差 + 归一化(Pre-Norm 风格)
x = x + self.dropout(self.attn(self.norm1(x), mask))
x = x + self.dropout(self.ff(self.norm2(x)))
return x
# 堆叠编码器
class TransformerEncoder(nn.Module):
def __init__(self, d_model, n_heads,
d_ff, n_layers, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList([
TransformerBlock(d_model, n_heads,
d_ff, dropout)
for _ in range(n_layers)
])
def forward(self, x, mask=None):
for layer in self.layers:
x = layer(x, mask)
return x
使用内置 Transformer
# PyTorch 内置实现(生产直接用)
encoder_layer = nn.TransformerEncoderLayer(
d_model=512, nhead=8, dim_feedforward=2048,
dropout=0.1, batch_first=True)
transformer = nn.TransformerEncoder(
encoder_layer, num_layers=6)
x = torch.rand(32, 20, 512) # [batch, seq, d]
out = transformer(x)
# 完整 Transformer(编码器+解码器)
model = nn.Transformer(
d_model=512, nhead=8,
num_encoder_layers=6, num_decoder_layers=6)
Transformer 应用
# 1. 文本分类(用编码器)
# 2. 机器翻译(编码器-解码器)
# 3. 语言模型/GPT(解码器)
# 4. 预训练 BERT/BERT 类模型
# 5. 视觉 Transformer(ViT)
# 现代实践: 直接用 HuggingFace transformers
from transformers import BertModel, GPT2LMHeadModel
Transformer 是深度学习近年的最重要架构。掌握注意力机制与编码器结构,就理解了 GPT/BERT 等大模型的骨架。
PyTorch Transformer
PyTorch Transformer
本节介绍 PyTorch 内置 Transformer 模块的完整用法,包括编码器、解码器与训练文本分类任务。
内置模块
import torch.nn as nn
# 单层
encoder_layer = nn.TransformerEncoderLayer(
d_model=512, # 特征维度
nhead=8, # 注意力头数
dim_feedforward=2048, # 前馈层维度
dropout=0.1,
activation="relu",
batch_first=True, # 输入 [batch, seq, dim]
)
# 多层编码器
encoder = nn.TransformerEncoder(
encoder_layer, num_layers=6)
# 解码器
decoder_layer = nn.TransformerDecoderLayer(
d_model=512, nhead=8, batch_first=True)
decoder = nn.TransformerDecoder(
decoder_layer, num_layers=6)
输入输出形状
# batch_first=True 时:
# 输入: [batch, seq_len, d_model]
# 输出: [batch, seq_len, d_model]
batch, seq, d_model = 32, 20, 512
x = torch.rand(batch, seq, d_model)
encoder_out = encoder(x)
print(encoder_out.shape) # [32, 20, 512]
# 带 mask(填充位置)
src_key_padding_mask = torch.zeros(batch, seq,
dtype=torch.bool)
src_key_padding_mask[:, 15:] = True # 15 之后是填充
encoder_out = encoder(x,
src_key_padding_mask=src_key_padding_mask)
文本分类模型
class TextClassifier(nn.Module):
def __init__(self, vocab_size, d_model=128,
n_heads=4, n_layers=2, n_classes=2):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, d_model, padding_idx=0)
self.pos_encoding = PositionalEncoding(d_model)
layer = nn.TransformerEncoderLayer(
d_model=d_model, nhead=n_heads,
dim_feedforward=512, batch_first=True)
self.encoder = nn.TransformerEncoder(
layer, num_layers=n_layers)
self.fc = nn.Linear(d_model, n_classes)
def forward(self, x, mask=None):
# x: [batch, seq] 词索引
emb = self.embedding(x) # [b, s, d]
emb = self.pos_encoding(emb)
out = self.encoder(emb, mask) # [b, s, d]
pooled = out.mean(dim=1) # 平均池化
return self.fc(pooled)
训练分类器
model = TextClassifier(vocab_size=10000)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
model.train()
for x, y in train_loader:
# 构造 padding mask
pad_mask = (x == 0) # [batch, seq]
optimizer.zero_grad()
out = model(x, pad_mask)
loss = criterion(out, y)
loss.backward()
optimizer.step()
注意力掩码(Mask)类型
# 1. padding mask: 忽略填充位置
src_key_padding_mask = (tokens == 0)
# 2. 因果 mask: 只允许看前面(生成任务)
# 用 attn_mask 下三角
seq = 20
causal_mask = torch.tril(torch.ones(seq, seq)
).bool()
# 3. 组合使用
# 解码器专用参数:
# tgt_mask: 目标序列因果掩码
# memory_key_padding_mask: 编码器输出的填充掩码
推理与生成
model.eval()
with torch.no_grad():
x = torch.tensor([[12, 45, 3, 0, 0]])
mask = (x == 0)
out = model(x, mask)
pred = out.argmax(dim=1)
print(pred.item())
性能与技巧
# 1. 位置编码后加 dropout
# 2. d_model 与 nhead 匹配(d_model % nhead == 0)
# 3. 训练加 warmup + 学习率调度
# 4. 长序列用 Flash Attention(2.0+ 自动)
# 5. 内存优化: gradient checkpointing
# Flash Attention(自动使用,无需改动)
torch.backends.cuda.enable_flash_sdp(True)
与 HuggingFace 对比
# PyTorch 内置: 教学/轻量场景
# HuggingFace: 生产级预训练模型
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-chinese", num_labels=2)
PyTorch 内置 Transformer 模块让实现变得简单:定义层、堆叠、加 mask、训练。要更强大就用预训练模型(HuggingFace)。
PyTorch torch.optim
PyTorch torch.optim
torch.optim 提供各类优化器与学习率调度器,是训练循环中"更新参数"的核心组件。
优化器概览
import torch.optim as optim
# 梯度下降家族
optim.SGD(params, lr=0.01, momentum=0.9,
weight_decay=1e-4)
optim.Adagrad / optim.RMSprop
# 自适应学习率(默认首选)
optim.Adam(params, lr=1e-3, betas=(0.9, 0.999),
weight_decay=0)
optim.AdamW(params, lr=1e-3, weight_decay=0.01)
# AdamW: 解耦权重衰减,Transformer 训练标准选择
# 其他
optim.Adamax / optim.Adadelta / optim.NAdam
optim.LBFGS # 小批量二阶优化
基本使用
model = nn.Linear(10, 2)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 训练循环中的标准用法
for epoch in range(10):
for x, y in loader:
optimizer.zero_grad() # 1. 清零梯度
loss = criterion(model(x), y)
loss.backward() # 2. 反传
optimizer.step() # 3. 更新
# 等价: model.zero_grad() 清零所有参数梯度
按组设置不同学习率
# 不同层不同 lr(如微调预训练模型)
optimizer = optim.Adam([
{"params": model.backbone.parameters(),
"lr": 1e-5}, # 骨干: 小 lr
{"params": model.head.parameters(),
"lr": 1e-3}, # 分类头: 正常 lr
], lr=1e-3)
# 冻结部分参数(不更新)
for p in model.backbone.parameters():
p.requires_grad = False
梯度裁剪
# 防止梯度爆炸(RNN/Transformer 常见)
optimizer.zero_grad()
loss.backward()
# 按范数裁剪
torch.nn.utils.clip_grad_norm_(
model.parameters(), max_norm=1.0)
# 按值裁剪
torch.nn.utils.clip_grad_value_(
model.parameters(), clip_value=0.5)
optimizer.step()
学习率调度器
# 创建调度器
scheduler = optim.lr_scheduler.StepLR(
optimizer, step_size=30, gamma=0.1)
# 每 30 个 epoch, lr *= 0.1
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=100)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode="min", patience=5, factor=0.5)
# 验证损失不降时自动降低 lr
# 训练循环中的调用
for epoch in range(100):
train_one_epoch()
val_loss = evaluate()
# 方式1: 按 epoch 调度
scheduler.step()
# 方式2: 按指标调度(ReduceLROnPlateau)
# scheduler.step(val_loss)
# 查看当前 lr
print(optimizer.param_groups[0]["lr"])
Warmup 调度
# Transformer 训练常用 warmup
def warmup_schedule(step, warmup=1000, d_model=512):
"""warmup 后按倒数平方根衰减"""
step = max(step, 1)
return (d_model ** -0.5) * min(step ** -0.5,
step * warmup ** -1.5)
# 或用 LambdaLR
scheduler = optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda step: warmup_schedule(step))
最佳实践
# 1. 默认起点: AdamW + weight_decay + Cosine
# 2. 小数据/简单任务: SGD + momentum 也不错
# 3. 微调: 小 lr + 分层学习率
# 4. RNN/Transformer: 梯度裁剪 + warmup
# 5. 验证集监控: ReduceLROnPlateau 最省心
# 6. lr 太高: 损失震荡/发散 → 调低 10 倍
# 7. lr 太低: 收敛慢 → 调高 10 倍
# 常用 lr 起点:
# Adam: 1e-3 ~ 3e-4
# SGD: 0.01 ~ 0.1
# 微调: 1e-5 ~ 5e-5
优化器决定"怎么更新",调度器决定"何时更新多快"。掌握 AdamW 与 Cosine/Plateau 组合,大多数任务都能收敛良好。
PyTorch torchvision
PyTorch torchvision
torchvision 是 PyTorch 的视觉工具包:数据集、预训练模型、图像变换、模型 zoo。它是计算机视觉开发的标配。
模块结构
import torchvision
torchvision.datasets # 内置数据集
torchvision.models # 预训练模型
torchvision.transforms # 图像变换
torchvision.io # 图像/视频 IO
torchvision.ops # 检测/分割算子
# 版本
print(torchvision.__version__)
内置数据集
import torchvision.datasets as datasets
import torchvision.transforms as transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
])
# 分类
datasets.MNIST / FashionMNIST
datasets.CIFAR10 / CIFAR100
datasets.SVHN
# 检测/分割
datasets.COCODetection(root, annFile)
datasets.VOCDetection(root, year)
# 其他
datasets.ImageFolder(root) # 自定义文件夹
datasets.ImageNet(root) # 需申请权限
# 用法
trainset = datasets.CIFAR10(
root="./data", train=True,
download=True, transform=transform)
预训练模型(models)
import torchvision.models as models
# 分类模型
models.resnet18(pretrained=True)
models.resnet50(pretrained=True)
models.vgg16(pretrained=True)
models.alexnet(pretrained=True)
models.efficientnet_b0(pretrained=True)
models.densenet121(pretrained=True)
models.mobilenet_v2(pretrained=True)
# 最新版使用 weights 参数(推荐)
models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
# 检测/分割
models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
models.detection.ssdlite320_mobilenet_v3_large(pretrained=True)
models.segmentation.deeplabv3_resnet50(pretrained=True)
# 注意: 预训练模型下载到 ~/.cache/torch/hub/
迁移学习
# 用预训练模型 + 新分类头
import torch.nn as nn
model = models.resnet18(pretrained=True)
# 冻结骨干
for param in model.parameters():
param.requires_grad = False
# 替换分类头(ImageNet 1000 类 → 自己 N 类)
num_features = model.fc.in_features # 512
model.fc = nn.Linear(num_features, 10)
# 只训练分类头
optimizer = optim.Adam(model.fc.parameters(), lr=1e-3)
# 微调(解冻部分层)
for param in model.layer4.parameters():
param.requires_grad = True
模型输入预处理
# 预训练模型要求: 224x224, ImageNet 归一化
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# 预测单张图片
from PIL import Image
img = Image.open("cat.jpg")
x = preprocess(img).unsqueeze(0) # [1, 3, 224, 224]
model.eval()
with torch.no_grad():
out = model(x)
prob = torch.softmax(out, dim=1)
top5 = prob.topk(5)
# ImageNet 类别标签
# 见 torchvision.models 的权重 meta
weights = models.ResNet50_Weights.IMAGENET1K_V2
labels = weights.meta["categories"]
transforms 增强
import torchvision.transforms as T
# 数据增强组合(训练)
train_tf = T.Compose([
T.RandomResizedCrop(224, scale=(0.08, 1.0)),
T.RandomHorizontalFlip(),
T.ColorJitter(0.4, 0.4, 0.4),
T.ToTensor(),
T.Normalize(mean, std),
])
# 高级增强(AutoAugment 等)
T.AutoAugment(T.AutoAugmentPolicy.IMAGENET)
T.RandAugment(num_ops=2, magnitude=9)
# 转换工具
T.ToPILImage() / T.ToTensor()
T.functional.resize(img, 256)
torchvision.ops
import torchvision.ops as ops
# 检测常用算子
ops.nms(boxes, scores, iou_threshold) # 非极大抑制
ops.roi_align(features, boxes, 7) # ROI 对齐
ops.box_iou(boxes1, boxes2) # IoU 计算
# 图像
ops.deform_conv2d / ops.misc
最佳实践
# 1. 数据集: 优先内置,ImageFolder 处理自定义
# 2. 模型: 优先预训练(迁移学习),少从零训
# 3. 预处理: 与预训练权重匹配(尺寸+归一化)
# 4. 增强: 训练集用,验证集只归一化
# 5. 硬件: 批量 GPU 训练考虑 num_workers/pin_memory
torchvision 让视觉开发"开箱即用":数据集、预训练模型、变换一站式。90% 的视觉任务用它+迁移学习即可起步。
PyTorch 模型部署
PyTorch 模型部署
模型训练完成后需要部署到生产环境。本节介绍 PyTorch 模型部署的主要方式:TorchScript、ONNX、TorchServe 与 FastAPI。
部署方式总览
# 1. TorchScript: PyTorch 原生序列化(跨语言 C++/Java)
# 2. ONNX: 开放格式,转换到其他推理引擎
# 3. TorchServe: 官方模型服务(生产级)
# 4. FastAPI: Python 服务(简单场景)
# 5. 移动端: iOS/Android(量化/轻量模型)
方式1: TorchScript 导出
import torch
# 训练好的模型
model = MyModel().eval()
# 方法A: tracing(推荐简单模型)
example = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example)
traced.save("model.pt")
# 方法B: scripting(需要控制流)
scripted = torch.jit.script(model)
scripted.save("model_script.pt")
# 加载推理
loaded = torch.jit.load("model.pt")
with torch.no_grad():
out = loaded(example)
方式2: ONNX 导出
import torch
model = MyModel().eval()
example = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model, # 模型
example, # 示例输入
"model.onnx", # 输出文件
export_params=True,
opset_version=17,
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={ # 动态 batch
"input": {0: "batch"},
"output": {0: "batch"},
},
)
# 验证 ONNX
import onnx
onnx.checker.check_model(onnx_model)
# 用 ONNX Runtime 推理(CPU/GPU)
import onnxruntime as ort
sess = ort.InferenceSession("model.onnx")
result = sess.run(["output"],
{"input": example.numpy()})
方式3: TorchServe 服务
# 1. 安装
pip install torchserve torch-model-archiver
# 2. 打包模型
torch-model-archiver --model-name mymodel \
--version 1.0 --serialized-file model.pt \
--handler handler.py --export-path model_store
# 3. 启动服务
torchserve --start --ncs \
--model-store model_store --models mymodel.mar
# 4. 调用 API
# POST http://localhost:8080/predictions/mymodel
# 返回 JSON 预测结果
# handler.py 定义预处理与后处理
# 支持: 自定义推理逻辑、批处理、指标监控
方式4: FastAPI 服务
# app.py
from fastapi import FastAPI
import torch
app = FastAPI()
model = torch.jit.load("model.pt")
model.eval()
@app.post("/predict")
def predict(data: dict):
x = torch.tensor(data["features"]).unsqueeze(0)
with torch.no_grad():
out = model(x)
return {"prediction": out.tolist()}
# 运行: uvicorn app:app --port 8000
# 测试: curl -X POST localhost:8000/predict -d '{"features":[...]}'
推理优化
# 1. 模型量化(减小体积,提升速度)
model.qconfig = torch.quantization.get_default_qconfig("fbgemm")
quantized = torch.quantization.quantize_dynamic(
model, {torch.nn.Linear}, dtype=torch.qint8)
# 2. 半精度(FP16)
model = model.half()
x = x.half()
# 3. 编译加速(2.0+)
model = torch.compile(model)
# 4. 批处理与缓存
# 5. 多线程推理
部署注意事项
# 1. 训练/推理模式: eval() + no_grad()
# 2. 输入预处理与训练时一致(归一化等)
# 3. 固定输入形状/动态轴明确
# 4. 版本管理: 模型版本与代码版本对应
# 5. 监控: 延迟、吞吐、漂移检测
# 6. 回滚: 保存历史模型,支持快速回退
# 部署前验证:
# - 本地推理结果与训练时一致
# - 边界输入(空、异常)不崩溃
# - 性能达标(延迟/吞吐)
部署是模型的"最后一公里"。简单场景用 FastAPI,生产级用 TorchServe,跨平台/异构用 ONNX。按需选择,先跑通再优化。
PyTorch 模型保存和加载
PyTorch 模型保存和加载
模型持久化是训练必备:保存检查点、中断恢复、模型分享。本节介绍 PyTorch 的保存加载机制。
核心概念
# state_dict: 模型参数与缓冲区的字典
# - model.state_dict() 获取
# - model.load_state_dict() 加载
# - 只含参数,不含结构
# 保存方式:
# 1. 只存参数(state_dict,推荐)
# 2. 存完整模型(不推荐,耦合代码)
# 3. 存检查点(参数+优化器+epoch)
方式1: 保存/加载参数
import torch
# 保存
torch.save(model.state_dict(), "model.pth")
# 加载
model = MyModel() # 先创建结构
model.load_state_dict(torch.load("model.pth"))
model.eval() # 推理模式
# 注意: load 后必须创建相同结构!
# 强烈建议: 保存结构信息(如 config.json)
方式2: 保存检查点(恢复训练)
# 保存: 模型 + 优化器 + 调度器 + epoch
checkpoint = {
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"best_acc": best_acc,
}
torch.save(checkpoint, "checkpoint.pth")
# 加载恢复
checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
epoch = checkpoint["epoch"]
best_acc = checkpoint["best_acc"]
# 继续训练
for e in range(epoch, total_epochs):
...
多 GPU 保存与加载
# DataParallel 保存(名字含 module.)
torch.save(model.module.state_dict(), "model.pth")
# 加载兼容
state_dict = torch.load("model.pth")
# 去除 module. 前缀
from collections import OrderedDict
new_sd = OrderedDict(
(k.replace("module.", ""), v)
for k, v in state_dict.items())
model.load_state_dict(new_sd)
设备迁移
# 保存时在 GPU,加载时在 CPU
state_dict = torch.load(
"model.pth",
map_location=torch.device("cpu"))
model.load_state_dict(state_dict)
# 加载到 GPU
state_dict = torch.load(
"model.pth",
map_location="cuda:0")
# 或统一处理
state_dict = torch.load(
"model.pth",
map_location=lambda storage, loc: storage)
# 加载后手动 model.to(device)
完整模型保存(不推荐)
# 保存完整模型(含结构)
torch.save(model, "model_full.pth")
# 加载(需要模型类可导入!)
model = torch.load("model_full.pth")
# 缺点: 代码耦合、版本升级易碎
# 推荐: 参数 + 结构配置分开存
HuggingFace 格式(生态)
# transformers 模型有自己的保存格式
model.save_pretrained("my_model/")
# 生成: config.json + model.safetensors
# 加载
from transformers import AutoModel
model = AutoModel.from_pretrained("my_model/")
# 上传分享
# from huggingface_hub import upload_folder
最佳实践
# 1. 每 epoch 保存一次检查点(或按 best 保存)
# 2. 保留 best 模型 + 最近检查点
# 3. 文件名带 epoch/acc: model_ep5_acc97.pth
# 4. 保存完整训练状态(可恢复)
# 5. 分享时附上模型结构说明/config
# 6. 大模型用 safetensors 格式(安全高效)
# 定期自动保存
best_acc = 0
for epoch in range(epochs):
acc = train_epoch()
if acc > best_acc:
best_acc = acc
torch.save(model.state_dict(), "best.pth")
torch.save(checkpoint, f"ckpt_{epoch}.pth")
保存加载三件套:state_dict 存参数、checkpoint 存训练状态、config 存结构。按场景选择,模型管理不再混乱。
PyTorch 图像分类
PyTorch 图像分类
图像分类是计算机视觉的基础任务。本节用一个完整案例(CIFAR10 + ResNet 迁移学习),演示工业级的图像分类流程。
完整流程
1. 数据准备(下载+增强+加载)
2. 模型构建(预训练+改分类头)
3. 训练(优化器+调度+监控)
4. 评估(准确率/混淆矩阵)
5. 推理(单张预测)
6. 保存(模型+配置)
1. 数据准备
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as T
# 增强(训练)与归一化(测试)
train_tf = T.Compose([
T.RandomCrop(32, padding=4),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
test_tf = T.Compose([
T.ToTensor(),
T.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
train_set = torchvision.datasets.CIFAR10(
"./data", train=True, transform=train_tf,
download=True)
test_set = torchvision.datasets.CIFAR10(
"./data", train=False, transform=test_tf,
download=True)
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=128, shuffle=True,
num_workers=4)
test_loader = torch.utils.data.DataLoader(
test_set, batch_size=256, num_workers=4)
2. 模型(ResNet 迁移)
import torchvision.models as models
# 预训练 ResNet18
model = models.resnet18(pretrained=True)
# 适配 CIFAR 小图(32x32): 修改第一层卷积
model.conv1 = nn.Conv2d(3, 64, kernel_size=3,
stride=1, padding=1, bias=False)
model.maxpool = nn.Identity() # 去掉 maxpool
# 替换分类头
model.fc = nn.Linear(model.fc.in_features, 10)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
3. 训练
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3,
weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=30)
best_acc = 0
for epoch in range(30):
# 训练
model.train()
total_loss = 0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
total_loss += loss.item()
# 验证
acc = evaluate(model, test_loader, device)
scheduler.step()
print(f"Epoch {epoch+1}: loss={total_loss/len(train_loader):.4f}, "
f"acc={acc:.4f}")
# 保存 best
if acc > best_acc:
best_acc = acc
torch.save(model.state_dict(), "best_cifar.pth")
4. 评估函数
def evaluate(model, loader, device):
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
pred = model(x).argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
# 每类准确率
def per_class_acc(model, loader, device, n_classes=10):
model.eval()
correct = torch.zeros(n_classes)
total = torch.zeros(n_classes)
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
pred = model(x).argmax(dim=1)
for c in range(n_classes):
mask = (y == c)
correct[c] += (pred[mask] == c).sum().item()
total[c] += mask.sum().item()
return correct / total.clamp(min=1)
# 混淆矩阵
# from sklearn.metrics import confusion_matrix
5. 推理
from PIL import Image
def predict_image(model, img_path, transform, device):
model.eval()
img = Image.open(img_path).convert("RGB")
x = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
out = model(x)
prob = torch.softmax(out, dim=1)
pred = out.argmax(dim=1).item()
return pred, prob[0][pred].item()
# CIFAR10 类别
classes = ["airplane", "automobile", "bird", "cat",
"deer", "dog", "frog", "horse", "ship", "truck"]
pred, conf = predict_image(model, "cat.png", test_tf, device)
print(f"预测: {classes[pred]} ({conf:.2%})")
6. 训练技巧提升准确率
# 基线: ResNet18 迁移,约 90%+
# 进一步提升:
# - 用 ResNet50/更强大模型
# - 增加训练轮数与数据增强强度
# - Label Smoothing
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
# - 混合精度训练(AMP)
# - 学习率 warmup
完整训练模板
# 上面代码可组合成 train.py 脚本
# 通用模板适用其他数据集:
# 1. 改数据加载(类别数/尺寸)
# 2. 改模型分类头(num_classes)
# 3. 训练循环不变
# 推荐工具: PyTorch Lightning 简化工程化
# import pytorch_lightning as pl
图像分类的工业流程:数据增强 + 预训练迁移 + 余弦调度 + best 保存。这套模板迁移到其他分类任务,只需改数据和类别数。
PyTorc 文本情感分析
PyTorch 文本情感分析
情感分析判断文本的情绪倾向(正面/负面)。本节用 PyTorch 实现完整的情感分类:数据准备、词嵌入、模型、训练与评估。
任务定义
# 输入: 评论文本
# 输出: 情感类别(正面/负面)
# 流程:
# 文本 → 分词 → 词索引 → Embedding → 序列模型 → 分类
1. 数据准备
# 使用 torchtext 的 IMDB 数据集
from torchtext.datasets import IMDB
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
tokenizer = get_tokenizer("basic_english")
# 构建词表
def yield_tokens(data_iter):
for label, text in data_iter:
yield tokenizer(text)
train_iter, test_iter = IMDB(split=("train", "test"))
vocab = build_vocab_from_iterator(
yield_tokens(train_iter),
specials=["", ""])
vocab.set_default_index(vocab[""])
print(f"词表大小: {len(vocab)}")
# 文本 → 索引序列
def text_to_indices(text):
return [vocab[token] for token in tokenizer(text)]
# 截断/填充到固定长度
MAX_LEN = 200
def process(text):
ids = text_to_indices(text)[:MAX_LEN]
ids += [vocab[""]] * (MAX_LEN - len(ids))
return torch.tensor(ids)
# 自定义 Dataset
from torch.utils.data import Dataset, DataLoader
class SentimentDataset(Dataset):
def __init__(self, data_iter):
self.samples = [(process(text), 1 if label == 2 else 0)
for label, text in data_iter]
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
return self.samples[idx]
train_ds = SentimentDataset(IMDB(split="train"))
test_ds = SentimentDataset(IMDB(split="test"))
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=256)
2. 模型定义
import torch.nn as nn
class SentimentModel(nn.Module):
def __init__(self, vocab_size, embed_dim=128,
hidden_dim=128, n_classes=2):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=1) # =1
self.lstm = nn.LSTM(
embed_dim, hidden_dim,
batch_first=True, bidirectional=True)
self.fc = nn.Linear(hidden_dim * 2, n_classes)
self.dropout = nn.Dropout(0.3)
def forward(self, x, lengths=None):
emb = self.embedding(x) # [b, s, d]
out, _ = self.lstm(emb) # [b, s, 2h]
last = out[:, -1, :] # 最后步
return self.fc(self.dropout(last))
model = SentimentModel(len(vocab))
print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")
3. 训练
import torch.optim as optim
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
def train_epoch(loader):
model.train()
total_loss = correct = total = 0
for x, y in loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(), 1.0) # 梯度裁剪
optimizer.step()
total_loss += loss.item()
correct += (out.argmax(1) == y).sum().item()
total += y.size(0)
return total_loss / len(loader), correct / total
def evaluate(loader):
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
correct += (model(x).argmax(1) == y).sum().item()
total += y.size(0)
return correct / total
for epoch in range(5):
loss, acc = train_epoch(train_loader)
test_acc = evaluate(test_loader)
print(f"Epoch {epoch+1}: loss={loss:.4f}, "
f"train_acc={acc:.4f}, test_acc={test_acc:.4f}")
4. 预测
def predict(text, model, max_len=200):
model.eval()
ids = text_to_indices(text)[:max_len]
ids += [vocab[""]] * (max_len - len(ids))
x = torch.tensor([ids]).to(device)
with torch.no_grad():
out = model(x)
prob = torch.softmax(out, dim=1)
pred = out.argmax(1).item()
label = "正面" if pred == 1 else "负面"
return label, prob[0][pred].item()
print(predict("This movie is fantastic and well made!"))
# 输出: 正面, 0.98
print(predict("Terrible acting and boring plot."))
# 输出: 负面, 0.95
5. 改进方向
# 1. 用预训练词向量(GloVe/FastText)初始化 embedding
# 2. 用 Transformer 替代 LSTM
# 3. 用 HuggingFace 预训练模型(BERT)微调
# → 效果最佳,数据少也强
# HuggingFace 方案(一行起):
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-uncased", num_labels=2)
关键点回顾
# 1. 文本 → 词索引 → Embedding
# 2. 序列模型(LSTM)聚合信息
# 3. 最后时间步 → 分类头
# 4. 梯度裁剪防止 RNN 梯度爆炸
# 5. 填充 padding 与掩码
# 工业级建议:
# 中文用 jieba 分词 + BERT 类模型
# 数据量小时迁移学习优于从零训练
情感分析是 NLP 入门任务:覆盖了分词、词表、嵌入、序列模型与训练评估全流程。掌握后,文本分类类任务(主题、垃圾邮件等)都是同一套路。
PyTorch Autograd
PyTorch Autograd
Autograd(自动求导)是 PyTorch 的核心引擎:自动计算梯度,让反向传播无需手写。本节深入理解其机制。
基本机制
import torch
# 1. 需要求导的张量
x = torch.tensor([2.0, 3.0], requires_grad=True)
# 2. 计算(构建计算图)
y = x ** 2 + 3 * x # 每个元素: x^2 + 3x
# 3. 反向传播(标量输出)
y.sum().backward()
# 4. 查看梯度 dy/dx = 2x + 3 → [7, 9]
print(x.grad) # tensor([7., 9.])
计算图
# Autograd 动态构建计算图:
# 每个运算记录: 输入、输出、导数函数
# backward 时沿图反向传播梯度
# 图结构(简化):
x → pow(2) → add → sum → backward
└→ mul(3) →┘
# 每个中间张量都有 .grad_fn(导数函数)
y.grad_fn # AddBackward
x.grad_fn # None(叶子节点)
# 叶子节点: requires_grad=True 的原始张量
# 只有叶子节点保存 .grad
梯度积累与清零
# backward 默认累加梯度!
x = torch.tensor([1.0], requires_grad=True)
y = x ** 2
y.backward()
print(x.grad) # [2.]
y.backward() # 再次反传
print(x.grad) # [4.] ← 累加了!
# 训练中必须清零
optimizer.zero_grad()
# 或
x.grad.zero_()
# 这正是训练循环里
# zero_grad 必须在 backward 之前的原因
控制求导
# 1. no_grad: 不构建图(推理/评估)
with torch.no_grad():
y = model(x) # 省内存,速度快
# 2. 局部禁用
torch.set_grad_enabled(False)
# 3. 原地修改 requires_grad
x.requires_grad_(False)
# 4. 分离张量(不再参与求导)
z = x.detach() # 从图中断开
# 典型用途: 对抗训练、target 网络
target = target_net(x).detach()
自定义导数
import torch
# 自定义 autograd 函数
class MyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return x.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
grad_x = grad_output.clone()
grad_x[x < 0] = 0 # ReLU 导数
return grad_x
x = torch.tensor([-1.0, 2.0], requires_grad=True)
y = MyReLU.apply(x)
y.sum().backward()
print(x.grad) # [0., 1.]
高阶导数
# 二阶导数(create_graph=True)
x = torch.tensor([2.0], requires_grad=True)
y = x ** 3 # y' = 3x^2, y'' = 6x
y.backward(create_graph=True)
first = x.grad.clone() # [12.]
x.grad.zero_()
first.sum().backward() # 对一阶导再反传
print(x.grad) # [12.] ← 二阶导 6*2
调试技巧
# 1. 梯度检查(数值对比)
x.requires_grad_(True)
y = f(x)
y.sum().backward()
# 手动计算近似梯度对比
# 2. 检查梯度是否为 NaN
for name, p in model.named_parameters():
if p.grad is not None and torch.isnan(p.grad).any():
print(f"NaN grad: {name}")
# 3. 梯度爆炸处理
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
# 4. requires_grad 检查
p.requires_grad # 是否参与训练
常见误区
# 1. 忘了 zero_grad → 梯度累加错误
# 2. 推理时忘了 no_grad → 内存爆炸
# 3. 原地操作(+=)破坏计算图
# → 用 x = x + 1 而非 x += 1
# 4. 非叶子节点取 .grad → None
# 5. detach 使用时机不清
# 内存: 梯度图在 backward 后自动释放
# 若需保留, backward(retain_graph=True)
Autograd 是 PyTorch 的"隐形引擎"。理解计算图、梯度积累与 no_grad,训练代码就不会出错,也能更好地调试梯度问题。
PyTorch GPU / CUDA 加速
PyTorch GPU / CUDA 加速
GPU 加速是深度学习的标配。本节介绍 CUDA 环境检查、张量/模型迁移、以及性能优化技巧。
检查 GPU 环境
import torch
print(torch.cuda.is_available()) # GPU 可用?
print(torch.cuda.device_count()) # GPU 数量
print(torch.cuda.get_device_name(0)) # GPU 名称
print(torch.cuda.current_device()) # 当前设备
print(torch.version.cuda) # CUDA 版本
# MPS(macOS Apple 芯片)
print(torch.backends.mps.is_available())
设备管理
# 统一设备
device = (
"cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available()
else "cpu"
)
print(f"使用设备: {device}")
# 张量迁移
x = torch.rand(3, 3)
x = x.to(device) # 通用写法
x = x.cuda() # 快捷(CUDA)
x = x.to("cuda:0") # 指定卡
# 模型迁移
model = model.to(device)
# 注意: 数据与模型必须在同一设备!
# 错误: model 在 cuda, 输入在 cpu → 报错
训练循环中的设备处理
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
for x, y in train_loader:
x = x.to(device) # 每个 batch 迁移
y = y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
单机多 GPU: DataParallel
# 简单并行(数据并行)
model = nn.DataParallel(model)
model = model.to("cuda")
# 原理: batch 分到各卡,同步梯度
# 注意: state_dict 键名带 module. 前缀
# 保存: torch.save(model.module.state_dict(), ...)
单机多 GPU: DistributedDataParallel(推荐)
import torch.distributed as dist
import torch.multiprocessing as mp
def main(rank, world_size):
# 初始化进程组
dist.init_process_group(
"nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
model = model.to(rank)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[rank])
# 数据并行分发
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, batch_size=64)
for epoch in range(10):
sampler.set_epoch(epoch) # 每轮打乱
train(model, loader)
dist.destroy_process_group()
if __name__ == "__main__":
mp.spawn(main, args=(torch.cuda.device_count(),),
nprocs=torch.cuda.device_count())
# 启动: torchrun --nproc_per_node=4 train.py
混合精度训练(AMP)
# FP16 训练: 省显存、加速、精度基本无损
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler() # 梯度缩放(防下溢)
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
with autocast(): # 混合精度前向
loss = criterion(model(x), y)
scaler.scale(loss).backward() # 缩放反传
scaler.step(optimizer) # 缩放更新
scaler.update() # 调整缩放因子
# PyTorch 2.x 简化写法
torch.set_amp_enabled(True)
# 或
with torch.autocast(device_type="cuda"):
loss = criterion(model(x), y)
性能优化清单
# 1. DataLoader: num_workers=4~8, pin_memory=True
# 2. 混合精度(AMP): 显存减半,速度 1.5-2x
# 3. 编译模式(2.0+): model = torch.compile(model)
# 4. 减少 CPU-GPU 传输: 批内合并操作
# 5. 避免同步: 少用 .item()/.numpy()(训练中)
# 6. 基准测试: torch.cuda.synchronize() 后计时
# 7. 梯度累积(显存不足时):
accum = 4 # 4 个 batch 更新一次
optimizer.zero_grad()
for i, (x, y) in enumerate(loader):
loss = criterion(model(x.to(device)), y.to(device))
loss = loss / accum # 平均
loss.backward()
if (i + 1) % accum == 0:
optimizer.step()
optimizer.zero_grad()
显存管理
# 查看显存
print(torch.cuda.memory_allocated() / 1024**3, "GB")
print(torch.cuda.memory_reserved() / 1024**3, "GB")
print(torch.cuda.max_memory_allocated() / 1024**3, "GB")
# 释放缓存
torch.cuda.empty_cache()
# 显存不足应对:
# 1. 减小 batch_size
# 2. 混合精度(减半)
# 3. gradient checkpointing(用时间换空间)
model = torch.utils.checkpoint.checkpoint_sequential(
layers, segments=4, input=x)
GPU 加速的核心:统一设备、AMP 混合精度、合理 DataLoader、必要时多卡。先单卡优化到极限,再上多卡扩展。
PyTorch 损失函数
PyTorch 损失函数
损失函数衡量模型输出与目标的差距,是优化方向。本节系统介绍 PyTorch 的常用损失函数及选择方法。
分类损失
import torch.nn as nn
# 多分类(最常用)
nn.CrossEntropyLoss()
# 内部 = LogSoftmax + NLLLoss
# 输入: [batch, num_classes] 原始得分
# 目标: [batch] 类别索引(0~n-1)
criterion = nn.CrossEntropyLoss()
loss = criterion(logits, labels)
# 多分类(手动 Softmax)
nn.NLLLoss() # 需要 log-probabilities
# 搭配: nn.LogSoftmax(dim=1)
# 二分类(推荐)
nn.BCEWithLogitsLoss()
# 输入: [batch, 1] 原始 logits
# 目标: [batch, 1] 0/1
# 内部含 sigmoid,数值更稳定
# 二分类(手动 sigmoid)
nn.BCELoss()
# 输入需要先经过 sigmoid
# 多标签分类
nn.MultiLabelSoftMarginLoss()
# 平滑标签(防过拟合)
nn.CrossEntropyLoss(label_smoothing=0.1)
回归损失
# 均方误差(默认选择)
nn.MSELoss()
# 大误差惩罚重,对离群点敏感
# 平均绝对误差
nn.L1Loss()
# 对离群点鲁棒,但梯度恒定
# SmoothL1 / Huber(推荐)
nn.SmoothL1Loss(beta=1.0)
# 小误差用 L2,大误差用 L1
# 回归任务常用,抗离群点
# 对比:
# 预测 3, 目标 1
# MSE: 4 L1: 2 SmoothL1(β=1): 1.5
其他常用损失
# KL 散度(分布差异/蒸馏)
nn.KLDivLoss(reduction="batchmean")
# 输入: log_softmax 输出
# 目标: softmax 输出
# 余弦相似(表征学习)
nn.CosineEmbeddingLoss(margin=0.3)
# 三元组损失(度量学习)
nn.TripletMarginLoss(margin=1.0)
# 交叉熵多标签变体
nn.MultiMarginLoss()
# 排序损失
nn.MarginRankingLoss(margin=0.5)
自定义损失
# 方式1: 普通函数
def focal_loss(logits, target, gamma=2.0, alpha=0.25):
"""Focal Loss(类别不平衡)"""
ce = nn.functional.cross_entropy(logits, target,
reduction="none")
pt = torch.exp(-ce)
return (alpha * (1 - pt) ** gamma * ce).mean()
# 方式2: 继承 Module
class FocalLoss(nn.Module):
def __init__(self, gamma=2.0, alpha=0.25):
super().__init__()
self.gamma = gamma
self.alpha = alpha
def forward(self, logits, target):
return focal_loss(logits, target,
self.gamma, self.alpha)
criterion = FocalLoss(gamma=2.0)
损失组合
# 多任务/多损失加权
loss = 0.7 * cls_loss + 0.3 * reg_loss
# 正则化损失(weight decay 已在优化器中)
# 或手动加 L2
l2 = sum(p.pow(2).sum() for p in model.parameters())
loss = task_loss + 1e-4 * l2
reduction 参数
# reduction 控制损失聚合方式
nn.MSELoss(reduction="mean") # 默认: 平均
nn.MSELoss(reduction="sum") # 求和
nn.MSELoss(reduction="none") # 逐样本(不聚合)
# none 的用途:
# - 逐样本加权
# - 掩码(忽略某些样本)
losses = nn.functional.mse_loss(pred, target,
reduction="none")
masked = losses * mask
loss = masked.sum() / mask.sum() # 只统计有效样本
选择指南
# 任务类型 → 推荐损失
# 多分类 → CrossEntropyLoss
# 二分类 → BCEWithLogitsLoss
# 多标签 → BCEWithLogitsLoss(逐标签)
# 回归 → MSELoss 或 SmoothL1Loss
# 类别不平衡 → FocalLoss / 加权 CrossEntropy
# 蒸馏 → KLDivLoss
# 度量学习 → TripletMarginLoss / Contrastive
# 加权交叉熵(不平衡):
class_weights = torch.tensor([1.0, 5.0]).to(device)
nn.CrossEntropyLoss(weight=class_weights)
损失函数是训练的"指南针":任务类型决定损失,不平衡加权重,特殊需求自定义。选对损失,训练事半功倍。
PyTorch 学习率调度器
PyTorch 学习率调度器
学习率(learning rate)是训练中最重要的超参数。调度器让学习率按计划变化,显著影响收敛质量。本节介绍常用调度策略。
为什么需要调度
# 训练初期: 大 lr 快速下降
# 训练后期: 小 lr 精细收敛
# 固定 lr 的问题:
# - 太大: 后期震荡不收敛
# - 太小: 前期收敛太慢
# 调度器: 自动按阶段调整 lr
常用调度器
import torch.optim as optim
# 1. StepLR: 每隔 step_size 乘 gamma
scheduler = optim.lr_scheduler.StepLR(
optimizer, step_size=30, gamma=0.1)
# 2. MultiStepLR: 在指定 epoch 衰减
scheduler = optim.lr_scheduler.MultiStepLR(
optimizer, milestones=[30, 60, 90], gamma=0.1)
# 3. ExponentialLR: 每轮乘 gamma
scheduler = optim.lr_scheduler.ExponentialLR(
optimizer, gamma=0.95)
# 4. CosineAnnealingLR: 余弦衰减(推荐)
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=100) # 半个余弦周期
# 5. ReduceLROnPlateau: 指标停滞时降 lr
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode="min", factor=0.5,
patience=5, min_lr=1e-6)
# 6. OneCycleLR: 先升后降(训练快)
scheduler = optim.lr_scheduler.OneCycleLR(
optimizer, max_lr=1e-2, steps_per_epoch=len(loader),
epochs=10)
# 7. LambdaLR: 自定义函数
scheduler = optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda=lambda e: 0.95 ** e)
Warmup 与组合
# 训练初期 warmup(Transformer 必备)
def warmup_cosine(epoch, warmup=5, total=100):
if epoch < warmup:
return (epoch + 1) / warmup # 线性上升
progress = (epoch - warmup) / (total - warmup)
return 0.5 * (1 + math.cos(math.pi * progress))
scheduler = optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda=warmup_cosine)
# 或组合调度(SequentialLR)
scheduler = optim.lr_scheduler.SequentialLR(
optimizer,
schedulers=[
optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda=lambda e: e / 5),
optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=95),
],
milestones=[5], # 第 5 个 epoch 切换
)
调度器在训练循环中的位置
# 按 epoch 调度
for epoch in range(100):
train_one_epoch()
val_loss = validate()
scheduler.step() # 每轮调用
# (StepLR/Cosine 等按 epoch)
print(f"lr: {optimizer.param_groups[0]['lr']:.6f}")
# 按指标调度(ReduceLROnPlateau)
for epoch in range(100):
train_one_epoch()
val_loss = validate()
scheduler.step(val_loss) # 传入指标!
查看与自定义 lr
# 查看当前 lr
current_lr = optimizer.param_groups[0]["lr"]
# 手动设置 lr(热启动/调试)
optimizer.param_groups[0]["lr"] = 1e-4
# 多参数组各自调度
optimizer = optim.Adam([
{"params": model.backbone.parameters(), "lr": 1e-4},
{"params": model.head.parameters(), "lr": 1e-3},
])
scheduler.step() # 各组按各自比例更新
# 可视化学习率曲线
import matplotlib.pyplot as plt
lrs = []
for epoch in range(100):
lrs.append(optimizer.param_groups[0]["lr"])
scheduler.step()
plt.plot(lrs)
调度策略选择
# 场景 → 推荐
# CNN 常规训练 → Cosine / MultiStep
# Transformer → warmup + Cosine(或 InverseSqrt)
# 无验证集指标 → 按 epoch 调度(Step/Cosine)
# 有验证集 → ReduceLROnPlateau(省心)
# 训练时间有限 → OneCycle
# 通用组合(多数任务表现好):
# AdamW + CosineAnnealingLR + 5% warmup
# 注意:
# 调度器.step() 在优化器.step() 之后
# 恢复训练时要把调度器状态也恢复
# (见模型保存章节 checkpoint)
学习率调度是"锦上添花"的关键:固定 lr 也能训,但调度器让收敛更稳更好。先试 Cosine,再按需调整。
PyTorch 迁移学习
PyTorch 迁移学习
迁移学习:把在大数据(如 ImageNet)上预训练的模型,应用到自己的小数据集任务。数据少时效果显著。本节介绍完整流程。
为什么迁移学习有效
# 预训练模型学到了通用特征:
# 低层: 边缘/纹理(通用)
# 高层: 物体部件/语义(较通用)
# 分类头: 任务特定(需替换)
# 迁移 = 保留通用特征 + 训练任务特定部分
# 优势: 数据少也能训,收敛快,效果好
两种策略
# 策略1: 特征提取(冻结骨干)
# 只训练分类头,数据少时首选
# 速度快、不易过拟合
# 策略2: 微调(Fine-tuning)
# 解冻部分/全部层继续训练
# 数据充足、任务差异大时效果更好
# 用较小学习率(1e-5 ~ 1e-4)
完整实现
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import torchvision.transforms as T
# 1. 加载预训练模型
model = models.resnet18(pretrained=True)
# 2. 替换分类头
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10) # 新任务 10 类
# 3. 策略选择
# 特征提取: 冻结全部骨干
for param in model.parameters():
param.requires_grad = False
# 只训练分类头
for param in model.fc.parameters():
param.requires_grad = True
# 微调: 解冻最后几个层
for param in model.layer4.parameters():
param.requires_grad = True
# 4. 优化器(只优化 requires_grad=True 的参数)
optimizer = optim.Adam(
filter(lambda p: p.requires_grad,
model.parameters()),
lr=1e-3)
# 或分层 lr
optimizer = optim.Adam([
{"params": model.fc.parameters(), "lr": 1e-3},
{"params": model.layer4.parameters(), "lr": 1e-4},
])
数据与预处理
# 预训练模型要求 224x224 + ImageNet 归一化
preprocess = T.Compose([
T.Resize(256),
T.CenterCrop(224),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# 训练增强
train_tf = T.Compose([
T.RandomResizedCrop(224),
T.RandomHorizontalFlip(),
T.RandomRotation(15),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
两阶段训练(推荐)
# 阶段1: 冻结骨干,只训分类头(几轮)
# 阶段2: 解冻骨干,小 lr 微调(多轮)
# 阶段1
for param in model.parameters():
param.requires_grad = False
model.fc.requires_grad_(True)
optimizer1 = optim.Adam(model.fc.parameters(), lr=1e-3)
train(model, optimizer1, epochs=5)
# 阶段2
for param in model.parameters():
param.requires_grad = True
optimizer2 = optim.Adam(model.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer2, T_max=15)
train(model, optimizer2, epochs=15, scheduler=scheduler)
常用预训练模型对比
# 模型 参数量 速度 准确率 适用
# ResNet18 11M 快 高 通用首选
# ResNet50 25M 中 更高 数据充足
# EfficientNet 灵活 中 高 均衡
# MobileNet 4M 很快 中 移动端
# ViT 86M+ 慢 最高 大数据
# 大模型在数据少时不一定更好,先试 ResNet18
其他领域的迁移
# NLP: HuggingFace 预训练模型
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-uncased", num_labels=2)
# 音频: torchaudio 预训练
# 目标检测: 预训练 Faster R-CNN
import torchvision.models.detection as det
model = det.fasterrcnn_resnet50_fpn(pretrained=True)
注意事项
# 1. 输入尺寸与归一化必须与预训练一致
# 2. 冻结层不计算梯度(省显存内存)
# 3. 微调 lr 要比从头训练小 10-100 倍
# 4. 数据与预训练领域差异大 → 多解冻层
# 5. 小数据集: 特征提取优先,防过拟合
# 6. 数据增强仍是必要的
# 验证层是否冻结
for name, p in model.named_parameters():
if p.requires_grad:
print(f"训练: {name}")
迁移学习是数据不足场景的"作弊码":预训练 + 换头 + 微调,小数据集也能达到不错效果。几乎所有视觉/NLP 任务都应从它起步。
PyTorch 批归一化
PyTorch 批归一化
批归一化(Batch Normalization, BN)加速训练、稳定收敛、提升泛化,是现代网络的标准组件。本节介绍其原理与用法。
为什么需要归一化
# 问题: 网络各层输入分布随训练变化
# (Internal Covariate Shift)
# → 收敛慢、对 lr 敏感、深层难训
# BN 解决: 每层输入归一化到均值0方差1
# 再学一个缩放和平移(γ、β)恢复表达能力
# → 收敛快、可用更大 lr、自带正则化
BN 用法
import torch.nn as nn
# 全连接后
nn.BatchNorm1d(num_features)
# 卷积后(通道数)
nn.BatchNorm2d(num_channels)
# 使用位置: 卷积/线性之后,激活之前(常见)
nn.Sequential(
nn.Linear(256, 128),
nn.BatchNorm1d(128),
nn.ReLU(),
nn.Linear(128, 10),
)
# 或激活之后(部分风格)
# 习惯上"Conv → BN → ReLU"最常见
训练/推理模式差异
# 训练: 用当前 batch 的均值方差
# 推理: 用训练时累积的 running_mean/var
model.train() # BN 用 batch 统计量
model.eval() # BN 用 running 统计量
# 必须正确切换!否则推理结果错误:
# - 忘了 eval(): 单样本推理时 BN 计算异常
# - 忘了 train(): 训练时 BN 不更新统计量
# 查看 BN 参数
for name, p in model.named_parameters():
if "bn" in name:
print(name, p.shape)
# 查看 running 统计量(缓冲区)
bn = model.bn1
bn.running_mean / bn.running_var / bn.num_batches_tracked
BN 与其他归一化的对比
# BatchNorm: 按 batch 归一化(图像 CNN 主流)
# LayerNorm: 按样本归一化(Transformer 主流)
# InstanceNorm: 按样本通道归一化(风格迁移)
# GroupNorm: 按组归一化(小 batch)
nn.BatchNorm1d/2d # [N, C, ...] 按通道
nn.LayerNorm(normalized_shape) # 按特征
nn.InstanceNorm2d(channels)
nn.GroupNorm(num_groups, channels)
# 选择:
# - 大 batch 视觉 → BatchNorm
# - 小 batch(目标检测) → GroupNorm
# - Transformer/NLP → LayerNorm
# - 风格迁移 → InstanceNorm
小 batch 问题
# batch 太小(≤8)时 BN 统计量噪声大
# 解决方案:
# 1. 用 GroupNorm(不依赖 batch)
# 2. 增大 batch(梯度累积)
# 3. 用 SyncBatchNorm(多卡同步统计)
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
# 推理 batch=1 时 BN 正常(用 running 统计量)
BN 与优化
# BN 允许更大学习率
optimizer = optim.Adam(model.parameters(), lr=3e-3)
# BN 层一般不用 weight_decay
# 或用 lr 分组排除 BN 参数
decay_params = [p for n, p in model.named_parameters()
if "bn" not in n and p.requires_grad]
no_decay_params = [p for n, p in model.named_parameters()
if "bn" in n and p.requires_grad]
optimizer = optim.AdamW([
{"params": decay_params, "weight_decay": 1e-4},
{"params": no_decay_params, "weight_decay": 0},
])
BN 的实践建议
# 1. 网络里"Conv/Linear → BN → ReLU"是标配
# 2. 务必正确切换 train()/eval()
# 3. 保存/加载模型时 BN 统计量随 state_dict 保存
# 4. 微调预训练模型: BN 参数也要微调
# 5. 模型转换(ONNX/TorchScript)前先 eval()
# 6. 测试时 BN 统计量固化(冻结 BN)
# 可用于 domain adaptation 等
# BN 的作用总结:
# 加速收敛 / 允许大 lr / 缓解过拟合 / 稳定深层网络
BN 是现代网络的"稳定器"。用对位置(Conv→BN→ReLU)、切对模式(train/eval),就能享受它带来的全部好处。
PyTorch LSTM / GRU
PyTorch LSTM / GRU
LSTM 与 GRU 是解决 RNN 长序列问题的经典结构。本节深入讲解其用法:参数、输入输出形状、双向与多层。
RNN 的梯度问题
# 普通 RNN 长序列时:
# 梯度连乘 → 指数级消失/爆炸
# → 记不住长距离信息
# LSTM: 门控机制(遗忘/输入/输出)
# 维护记忆单元 c,可选择保留/更新
# GRU: 简化版(更新/重置门),参数更少
# 实际使用: LSTM/GRU 几乎完全取代普通 RNN
LSTM 参数
import torch.nn as nn
lstm = nn.LSTM(
input_size=100, # 输入特征维度(词向量维度)
hidden_size=128, # 隐藏状态维度
num_layers=2, # 层数(堆叠)
bias=True,
batch_first=True, # 输入 [batch, seq, feat]
dropout=0.1, # 层间 dropout(num_layers>1)
bidirectional=False, # 双向
)
# 输出:
# out: [batch, seq, hidden*(1+bidirectional)]
# h_n: [num_layers*(1+bid), batch, hidden]
# c_n: 同 h_n
输入输出详解
# 输入 x: [batch, seq_len, input_size]
# 例: 32 个句子,每句 20 词,词向量 100 维
x = torch.randn(32, 20, 100)
# 初始状态(可选,默认 0)
h0 = torch.zeros(2, 32, 128) # 2 层
c0 = torch.zeros(2, 32, 128)
out, (h_n, c_n) = lstm(x, (h0, c0))
# out: [32, 20, 128] 每个时间步输出(最后一层)
# h_n: [2, 32, 128] 各层最后步的隐藏状态
# c_n: [2, 32, 128] 各层记忆单元
# 常见取法:
# 1. 最后时间步(分类): out[:, -1, :]
# 2. 最后一层最后步: h_n[-1]
# 3. 双向拼接: 取 h_n[-2] 和 h_n[-1] 拼接
GRU(简化版)
gru = nn.GRU(
input_size=100, hidden_size=128,
num_layers=2, batch_first=True)
out, h_n = gru(x)
# 无 c_n,参数更少(~25%),速度更快
# 小数据时 GRU 常比 LSTM 效果好(不易过拟合)
双向 LSTM
lstm = nn.LSTM(100, 128, batch_first=True,
bidirectional=True)
# 输出维度翻倍: 128 * 2 = 256
out, (h_n, c_n) = lstm(x)
# out: [32, 20, 256] (正向128 + 反向128)
# 最后步拼接:
# 正向最后步 = out[:, -1, :128]
# 反向最后步 = out[:, 0, 128:] (反向从头开始)
# 拼接: torch.cat([out[:, -1, :128],
# out[:, 0, 128:]], dim=1)
# 双向适用: 全文可见任务(分类/标注)
# 生成任务用单向(不能看未来)
完整分类模型
class BiLSTMClassifier(nn.Module):
def __init__(self, vocab_size, embed_dim=128,
hidden_dim=128, n_classes=2):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=0)
self.lstm = nn.LSTM(
embed_dim, hidden_dim,
batch_first=True, bidirectional=True)
self.fc = nn.Linear(hidden_dim * 2, n_classes)
self.dropout = nn.Dropout(0.3)
def forward(self, x):
emb = self.embedding(x) # [b, s, d]
out, _ = self.lstm(emb) # [b, s, 2h]
# 平均池化(比取最后步更稳)
pooled = out.mean(dim=1) # [b, 2h]
return self.fc(self.dropout(pooled))
model = BiLSTMClassifier(vocab_size=10000)
变长序列处理
from torch.nn.utils.rnn import (
pack_padded_sequence, pad_packed_sequence)
def lstm_forward(model, x, lengths):
# x: [batch, seq], lengths: 每句真实长度
emb = model.embedding(x)
# 打包(跳过填充计算,更高效)
packed = pack_padded_sequence(
emb, lengths.cpu(), batch_first=True,
enforce_sorted=False)
out_packed, _ = model.lstm(packed)
# 解包回 [batch, seq, hidden]
out, _ = pad_packed_sequence(
out_packed, batch_first=True)
return out
# 取真实最后步(非填充位置)
def last_timestep(out, lengths):
batch = out.size(0)
idx = (lengths - 1).to(out.device)
idx = idx.view(batch, 1, 1).expand(
batch, 1, out.size(-1))
return out.gather(1, idx).squeeze(1)
训练与技巧
# 1. 梯度裁剪(RNN 必备)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
# 2. dropout 设置在层间与输出前
# 3. 初始状态一般用零(学习/指定均可)
# 4. 时间序列预测: 输出未来多步用循环/教师强制
# 5. 长序列: 用 Transformer 替代(效果更好)
# 与 Transformer 对比:
# LSTM: 参数少、适合中短序列、时间序列友好
# Transformer: 长距离强、可并行、数据要求高
LSTM/GRU 是序列建模的经典。掌握输入输出形状、双向与打包处理,就能应对分类、预测、生成等各类序列任务。
PyTorch 词嵌入
PyTorch 词嵌入
词嵌入(Word Embedding)把词映射为稠密向量,是 NLP 模型的基础。本节介绍 nn.Embedding 与预训练词向量。
为什么需要嵌入
# 词 → 索引(离散) → 向量(稠密, 可学习)
# 独热编码的问题:
# 维度大(词表大小)、无语义信息
# 嵌入: 低维稠密向量,相似的词向量相近
# 例: "猫" 和 "狗" 的向量距离 < "猫" 和 "汽车"
nn.Embedding 基本用法
import torch.nn as nn
# 词表 10000, 向量维度 128
embedding = nn.Embedding(
num_embeddings=10000, # 词表大小
embedding_dim=128, # 向量维度
padding_idx=0, # 填充位(向量全0,不训练)
)
# 输入: 词索引 [batch, seq]
x = torch.tensor([[12, 45, 3, 0], # "我 爱 学习 [pad]"
[7, 8, 9, 10]])
out = embedding(x) # [batch, seq, 128]
print(out.shape) # [2, 4, 128]
# 查单词向量
vec = embedding(torch.tensor(12)) # "我" 的向量
# 训练中嵌入会自动更新(是可学习参数)
预训练词向量加载
# 用 GloVe/FastText 预训练向量初始化
import torch
def load_glove(path, vocab, dim=100):
"""把 GloVe 向量加载到 embedding"""
embedding = nn.Embedding(
len(vocab), dim, padding_idx=vocab[""])
# 读取 GloVe 文件
vectors = {}
with open(path) as f:
for line in f:
parts = line.strip().split()
vectors[parts[0]] = [float(v)
for v in parts[1:]]
# 初始化(未命中用随机)
import torch.nn.init as init
init.normal_(embedding.weight, std=0.1)
for word, idx in vocab.get_stoi().items():
if word in vectors:
embedding.weight.data[idx] = torch.tensor(
vectors[word])
# 填充位置零
if embedding.padding_idx is not None:
embedding.weight.data[embedding.padding_idx] = 0
return embedding
# 或用 torchtext 的 Vectors
from torchtext.vocab import GloVe
glove = GloVe(name="6B", dim=100)
vec = glove["cat"] # 获取 "cat" 的向量
嵌入的相似度计算
import torch.nn.functional as F
def similarity(word1, word2, embedding, vocab):
idx1 = vocab[word1]
idx2 = vocab[word2]
v1 = embedding(torch.tensor(idx1))
v2 = embedding(torch.tensor(idx2))
return F.cosine_similarity(v1, v2, dim=0).item()
# 找最近邻
def nearest(word, embedding, vocab, k=5):
idx = vocab[word]
vec = embedding(torch.tensor(idx))
# 计算与所有词的余弦相似度
all_vecs = embedding.weight
scores = F.cosine_similarity(
vec.unsqueeze(0), all_vecs, dim=1)
top = scores.topk(k + 1) # 排除自己
return [vocab.get_itos()[i] for i in top.indices[1:]]
嵌入的进阶用法
# 1. 微调: 训练时冻结或低 lr
# 数据少时冻结:
embedding.weight.requires_grad = False
# 2. 维度选择: 100/128/256/300
# 大词表+大数据 → 高维;否则低维防过拟合
# 3. 特殊 token
# : 填充,向量恒 0
# : 未知词
# /: 序列开始/结束(生成任务)
# /: BERT 风格
# 4. 位置编码: 与嵌入相加注入位置
# (见 Transformer 章节 PositionalEncoding)
词嵌入 vs 上下文表示
# 静态嵌入(GloVe/Word2Vec):
# 一个词一个向量,多义词无法区分
# "bank" = 银行/河岸 同一向量
# 上下文嵌入(BERT/ELMo):
# 向量随上下文变化
# 现代 NLP 主流(Transformer 时代)
# PyTorch 生态:
# 静态: nn.Embedding + GloVe
# 上下文: HuggingFace BERT 等
# 从零训练小型任务: 静态嵌入足够
# 复杂语义任务: 用预训练上下文模型
嵌入可视化
# t-SNE/PCA 降维可视化
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
words = ["cat", "dog", "apple", "banana", "run", "walk"]
vecs = torch.stack([embedding(torch.tensor(vocab[w]))
for w in words]).numpy()
tsne = TSNE(n_components=2, random_state=0)
coords = tsne.fit_transform(vecs)
for i, word in enumerate(words):
plt.scatter(coords[i, 0], coords[i, 1])
plt.annotate(word, (coords[i, 0], coords[i, 1]))
plt.show()
词嵌入是 NLP 的"入口"。理解 nn.Embedding 用法、预训练向量加载与相似度计算,就能为任何文本模型打好输入基础。
PyTorch 生成对抗网络
PyTorch 生成对抗网络
生成对抗网络(GAN)由生成器与判别器对抗训练,能生成逼真的新样本。本节实现 DCGAN 生成手写数字。
GAN 原理
# 两个网络的博弈:
# 生成器 G: 噪声 → 假样本(试图骗过判别器)
# 判别器 D: 判断输入是真是假(试图识破生成器)
# 训练目标:
# D: 最大化区分真假的能力
# G: 最小化 D 识别出假的概率
# 达到平衡时: G 生成的样本足够逼真
# D 无法分辨(概率 0.5)
损失函数
# 判别器损失: 真实样本判真 + 假样本判假
d_loss = BCE(D(real), 1) + BCE(D(fake), 0)
# 生成器损失: 假样本骗过判别器
g_loss = BCE(D(fake), 1)
# 循环: 先训 D,再训 G(交替)
生成器(反卷积)
import torch.nn as nn
class Generator(nn.Module):
def __init__(self, z_dim=100, img_channels=1):
super().__init__()
self.model = nn.Sequential(
# 输入: [batch, z_dim, 1, 1]
nn.ConvTranspose2d(z_dim, 256, 4, 1, 0),
nn.BatchNorm2d(256), nn.ReLU(),
nn.ConvTranspose2d(256, 128, 4, 2, 1),
nn.BatchNorm2d(128), nn.ReLU(),
nn.ConvTranspose2d(128, 64, 4, 2, 1),
nn.BatchNorm2d(64), nn.ReLU(),
nn.ConvTranspose2d(64, img_channels, 4, 2, 1),
nn.Tanh(), # 输出 [-1, 1]
)
def forward(self, z):
return self.model(z.unsqueeze(2).unsqueeze(3))
# [batch, z_dim] → [batch, 1, 28, 28]
G = Generator()
判别器(卷积)
class Discriminator(nn.Module):
def __init__(self, img_channels=1):
super().__init__()
self.model = nn.Sequential(
nn.Conv2d(img_channels, 64, 4, 2, 1),
nn.LeakyReLU(0.2),
nn.Conv2d(64, 128, 4, 2, 1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
nn.Conv2d(128, 256, 4, 2, 1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.Conv2d(256, 1, 4, 1, 0), # [b,1,1,1]
nn.Sigmoid(),
)
def forward(self, x):
return self.model(x).view(-1, 1) # [batch, 1]
D = Discriminator()
训练循环
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as T
device = "cuda" if torch.cuda.is_available() else "cpu"
G = G.to(device); D = D.to(device)
# 数据: 图像归一化到 [-1, 1](与 Tanh 匹配)
transform = T.Compose([
T.ToTensor(),
T.Normalize((0.5,), (0.5,)),
])
mnist = datasets.MNIST("./data", train=True,
transform=transform, download=True)
loader = torch.utils.data.DataLoader(
mnist, batch_size=128, shuffle=True)
criterion = nn.BCELoss()
g_opt = optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
d_opt = optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
Z_DIM = 100
for epoch in range(30):
for real, _ in loader:
batch = real.size(0)
real = real.to(device)
labels_real = torch.ones(batch, 1).to(device)
labels_fake = torch.zeros(batch, 1).to(device)
# 训练判别器
d_opt.zero_grad()
z = torch.randn(batch, Z_DIM).to(device)
fake = G(z)
d_real = D(real)
d_fake = D(fake.detach()) # 不更新 G
d_loss = (criterion(d_real, labels_real)
+ criterion(d_fake, labels_fake)) / 2
d_loss.backward()
d_opt.step()
# 训练生成器
g_opt.zero_grad()
fake = G(z)
g_loss = criterion(D(fake), labels_real) # 骗过 D
g_loss.backward()
g_opt.step()
print(f"Epoch {epoch+1}: D={d_loss.item():.4f}, "
f"G={g_loss.item():.4f}")
生成新样本
import matplotlib.pyplot as plt
G.eval()
with torch.no_grad():
z = torch.randn(16, Z_DIM).to(device)
samples = G(z).cpu()
# 可视化
fig, axes = plt.subplots(4, 4, figsize=(6, 6))
for i, ax in enumerate(axes.flat):
img = (samples[i].squeeze() + 1) / 2 # 转 [0,1]
ax.imshow(img, cmap="gray")
ax.axis("off")
plt.show()
# 保存生成器(后续复用)
torch.save(G.state_dict(), "generator.pth")
训练技巧
# 1. 标签平滑: 用 0.9 替代 1(更稳)
# 2. lr 用 2e-4, betas=(0.5, 0.999)(DCGAN 配置)
# 3. 判别器先于生成器更新
# 4. 生成器用 LeakyReLU,不用 ReLU
# 5. 输出层 Tanh, 输入归一化到 [-1, 1]
# 6. 监控 D 的 loss: 稳定在 0.5 附近 = 平衡
# 7. 模式坍塌(生成样本单一): 调大噪声维度/调 lr
# 现代 GAN 变体: WGAN、LSGAN、StyleGAN
# 用对抗生成图像/音频/3D 内容(生成式 AI 基础)
GAN 让机器学会"创造"。理解生成器/判别器博弈与 DCGAN 结构,就能在此基础上探索各种生成模型。
PyTorch 自编码器
PyTorch 自编码器
自编码器(Autoencoder)是无监督学习模型:学习数据的压缩表示。用于降维、去噪、异常检测与特征学习。
原理
# 结构: 编码器 + 解码器
# 输入 → 编码器 → 低维表示(瓶颈) → 解码器 → 重建输入
# 训练目标: 重建误差最小
# loss = MSE(输出, 输入)
# 价值: 瓶颈层学到数据的核心特征(压缩表示)
# 应用: 降维 / 去噪 / 异常检测 / 预训练
基础自编码器
import torch
import torch.nn as nn
import torch.optim as optim
class Autoencoder(nn.Module):
def __init__(self, input_dim=784, latent_dim=32):
super().__init__()
# 编码器
self.encoder = nn.Sequential(
nn.Linear(input_dim, 256),
nn.ReLU(),
nn.Linear(256, latent_dim),
)
# 解码器
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 256),
nn.ReLU(),
nn.Linear(256, input_dim),
nn.Sigmoid(), # 输出 [0,1] 匹配图像
)
def forward(self, x):
z = self.encoder(x)
return self.decoder(z)
def encode(self, x):
return self.encoder(x) # 获取压缩表示
model = Autoencoder()
训练
# 数据: MNIST 展平
import torchvision.datasets as datasets
import torchvision.transforms as T
transform = T.Compose([T.ToTensor()]) # [0,1]
mnist = datasets.MNIST("./data", train=True,
transform=transform, download=True)
loader = torch.utils.data.DataLoader(
mnist, batch_size=128, shuffle=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
total_loss = 0
for x, _ in loader:
x = x.view(-1, 784).to(device) # 展平
optimizer.zero_grad()
recon = model(x)
loss = criterion(recon, x) # 重建误差
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Epoch {epoch+1}: loss={total_loss/len(loader):.4f}")
重建与可视化
import matplotlib.pyplot as plt
model.eval()
with torch.no_grad():
x, _ = next(iter(loader))
x = x.view(-1, 784).to(device)
recon = model(x)
# 对比原图与重建
fig, axes = plt.subplots(2, 8, figsize=(12, 3))
for i in range(8):
axes[0, i].imshow(x[i].view(28, 28).cpu(), cmap="gray")
axes[1, i].imshow(recon[i].view(28, 28).cpu(), cmap="gray")
axes[0, i].axis("off"); axes[1, i].axis("off")
plt.show()
降噪自编码器
class DenoiseAE(nn.Module):
def __init__(self, input_dim=784, latent_dim=32):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, 256), nn.ReLU(),
nn.Linear(256, latent_dim))
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 256), nn.ReLU(),
nn.Linear(256, input_dim), nn.Sigmoid())
def forward(self, x):
return self.decoder(self.encoder(x))
# 训练: 加噪输入, 重建原图
noise_factor = 0.3
for x, _ in loader:
x = x.view(-1, 784).to(device)
noisy = x + noise_factor * torch.randn_like(x)
noisy = noisy.clamp(0, 1)
recon = model(noisy)
loss = criterion(recon, x) # 与"干净"图比
...
卷积自编码器(图像更好)
class ConvAE(nn.Module):
def __init__(self):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 16, 3, 2, 1), # 28→14
nn.ReLU(),
nn.Conv2d(16, 32, 3, 2, 1), # 14→7
nn.ReLU(),
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(32, 16, 3, 2, 1, 1), # 7→14
nn.ReLU(),
nn.ConvTranspose2d(16, 1, 3, 2, 1, 1), # 14→28
nn.Sigmoid(),
)
def forward(self, x):
z = self.encoder(x)
return self.decoder(z)
# 输入保持 [batch, 1, 28, 28] 无需展平
变分自编码器(VAE)
# VAE: 学习分布的生成模型(可采样生成新样本)
import torch.nn.functional as F
class VAE(nn.Module):
def __init__(self, input_dim=784, latent_dim=32):
super().__init__()
self.fc_mu = nn.Linear(input_dim, latent_dim)
self.fc_logvar = nn.Linear(input_dim, latent_dim)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 256), nn.ReLU(),
nn.Linear(256, input_dim), nn.Sigmoid())
def encode(self, x):
h = F.relu(self.encoder_shared(x))
return self.fc_mu(h), self.fc_logvar(h)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + eps * std # 重参数化
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
return self.decoder(z), mu, logvar
# VAE 损失 = 重建 + KL 散度
def vae_loss(recon, x, mu, logvar):
recon_loss = F.binary_cross_entropy(
recon, x, reduction="sum")
kl_loss = -0.5 * torch.sum(
1 + logvar - mu.pow(2) - logvar.exp())
return recon_loss + kl_loss
# 生成新样本
with torch.no_grad():
z = torch.randn(16, 32).to(device)
samples = model.decoder(z)
应用场景
# 1. 降维: 高维数据 → 低维表示(可视化/加速)
# 2. 去噪: 图像/音频去噪
# 3. 异常检测: 重建误差大的样本 = 异常
# 正常样本重建好,异常样本重建差
# 4. 特征提取: 编码器输出作为下游任务特征
# 5. 生成(VAE): 采样生成新样本
# 异常检测示例
def anomaly_score(model, x, threshold):
recon = model(x)
error = ((recon - x) ** 2).mean(dim=1) # 重建误差
return (error > threshold).float() # 1=异常
自编码器是"无监督学习"的代表:不依赖标签学习数据结构。掌握 AE 与 VAE,就理解了表示学习与生成模型的入门。
PyTorch 模型评估与调试
PyTorch 模型评估与调试
训练只是开始,评估与调试决定模型质量。本节介绍评估指标、过拟合诊断与常见训练问题的排查方法。
评估指标
# 分类
accuracy # 准确率
precision # 精确率(预测为正中真为正)
recall # 召回率(正样本中被预测出的)
f1_score # 调和平均
confusion_matrix # 混淆矩阵
AUC/ROC # 二分类阈值无关指标
# 回归
MSE / MAE / RMSE
R2(决定系数)
# 序列/生成
BLEU / ROUGE(文本)
IoU / mAP(检测分割)
# PyTorch 计算
def accuracy(pred, target):
return (pred.argmax(1) == target).float().mean()
# sklearn 补充
from sklearn.metrics import (precision_score,
recall_score, f1_score, confusion_matrix)
评估流程
def evaluate_full(model, loader, device):
"""返回准确率与混淆矩阵"""
model.eval()
all_pred, all_true = [], []
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
pred = model(x).argmax(1)
all_pred.extend(pred.cpu().tolist())
all_true.extend(y.cpu().tolist())
from sklearn.metrics import (accuracy_score,
classification_report, confusion_matrix)
print(f"准确率: {accuracy_score(all_true, all_pred):.4f}")
print(classification_report(all_true, all_pred))
print(confusion_matrix(all_true, all_pred))
过拟合诊断
# 现象: 训练损失下降,验证损失上升
# 训练准确率 >> 验证准确率
# 应对:
# 1. 数据增强
# 2. 正则化: Dropout / weight_decay / 早停
# 3. 更小模型 / 减少层数
# 4. 更多数据 / 数据增强更强
# 5. Label Smoothing
# 6. 降低训练轮数
# 欠拟合: 训练损失也高
# → 加大模型 / 更多 epoch / 调大 lr
# 早停实现
best_val = float("inf")
patience = 5
wait = 0
for epoch in range(100):
train_loss = train_epoch()
val_loss = validate()
if val_loss < best_val:
best_val = val_loss
torch.save(model.state_dict(), "best.pth")
wait = 0
else:
wait += 1
if wait >= patience:
print(f"早停于 epoch {epoch}")
break
训练曲线可视化
import matplotlib.pyplot as plt
# 记录训练/验证损失与准确率
history = {"train_loss": [], "val_loss": [],
"train_acc": [], "val_acc": []}
# 训练中追加...
# 画图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
ax1.plot(history["train_loss"], label="train")
ax1.plot(history["val_loss"], label="val")
ax1.set_title("Loss"); ax1.legend()
ax2.plot(history["train_acc"], label="train")
ax2.plot(history["val_acc"], label="val")
ax2.set_title("Accuracy"); ax2.legend()
plt.show()
# 诊断:
# - 两线都高: 过拟合(早停/正则)
# - 两线分离大: 过拟合
# - 都低且不降: 欠拟合/学习率问题
# - 震荡: lr 太大
常见问题排查
# 1. Loss 为 NaN
# - 梯度爆炸 → 梯度裁剪 / 降低 lr
# - 数据含 NaN → 清洗数据
# - 除零 → 加 epsilon
# - lr 太大 → 调小
# 2. Loss 不降
# - lr 太小/太大 → 调 lr(10 倍试)
# - 数据问题(标签错误/未归一化)
# - 模型问题(激活/初始化)
# 3. 训练快但验证差(过拟合)
# - 增强 / Dropout / 早停 / 更小模型
# 4. 训练慢
# - GPU? 混合精度? num_workers?
# - 小 batch 试试
# 5. 结果随机波动
# - 固定种子(可复现)
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
# 或用 Lightning 的 seed_everything
调试工具
# 1. 梯度检查
def check_gradients(model):
for name, p in model.named_parameters():
if p.grad is not None:
if torch.isnan(p.grad).any():
print(f"NaN grad: {name}")
norm = p.grad.norm().item()
if norm > 100:
print(f"大梯度: {name} {norm:.1f}")
# 2. TensorBoard 可视化
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("runs/exp1")
writer.add_scalar("Loss/train", loss, epoch)
writer.add_scalar("Acc/val", acc, epoch)
writer.add_graph(model, sample_input)
# 运行: tensorboard --logdir runs
# 3. 输入输出检查
# 在 forward 中打印形状
# 或 hook 查看中间层
def hook_fn(module, input, output):
print(module.__class__.__name__,
output.shape)
model.layer1.register_forward_hook(hook_fn)
评估最佳实践
# 1. 训练/验证/测试三分离(测试集只看一次)
# 2. 用 best 模型(而非最后一轮)评估
# 3. 报告多指标,不只看准确率
# 4. 类别不平衡看 precision/recall/f1
# 5. 可视化错误样本(找模式)
# 6. 设置随机种子,实验可复现
# 7. 记录实验配置(超参/数据/代码版本)
# 简单实验管理
# - 手动: 文件名带参数(exp_lr1e3_b128)
# - 工具: Weights & Biases / MLflow
评估与调试是"训练的艺术":曲线诊断问题、指标衡量质量、工具定位根因。建立这套方法论,模型质量稳步提升。
PyTorch torchtext
PyTorch torchtext
torchtext 是 PyTorch 的文本处理工具包:分词、词表构建、数据集加载。让文本数据管线标准化。
torchtext 概览
import torchtext
torchtext.data.utils # 工具(get_tokenizer 等)
torchtext.datasets # 数据集(IMDB/AG_NEWS 等)
torchtext.vocab # 词表(build_vocab_from_iterator)
torchtext.data.functional # 函数式接口
# 注意: torchtext 0.17+ API 大改,以文档为准
分词
from torchtext.data.utils import get_tokenizer
# 英文分词
tokenizer = get_tokenizer("basic_english")
tokens = tokenizer("I love PyTorch!")
# ["i", "love", "pytorch"]
# 其他分词器
get_tokenizer("spacy") # 需 spacy
get_tokenizer("moses")
get_tokenizer("char") # 按字符
# 中文(需 jieba)
import jieba
def zh_tokenizer(text):
return list(jieba.cut(text))
构建词表
from torchtext.vocab import build_vocab_from_iterator
def yield_tokens(data_iter):
for label, text in data_iter:
yield tokenizer(text)
train_iter = IMDB(split="train")
vocab = build_vocab_from_iterator(
yield_tokens(train_iter),
specials=["", ""], # 特殊 token
min_freq=2, # 低频词丢弃
max_tokens=50000, # 词表上限
)
vocab.set_default_index(vocab[""]) # 未知词索引
print(len(vocab))
# 词表操作
vocab["hello"] # 词 → 索引
vocab.get_itos()[5] # 索引 → 词
vocab.get_stoi() # 词 → 索引映射
文本到张量
# 文本 → 索引序列
def text_to_tensor(text, max_len=200):
tokens = tokenizer(text)
ids = [vocab[t] for t in tokens][:max_len]
# 填充到固定长度
ids += [vocab[""]] * (max_len - len(ids))
return torch.tensor(ids, dtype=torch.long)
# 或动态 batch(用 padding 函数)
from torch.nn.utils.rnn import pad_sequence
def collate_batch(batch):
texts = [text_to_tensor(text) for _, text in batch]
labels = torch.tensor([1 if l == 2 else 0
for l, _ in batch])
padded = pad_sequence(texts, batch_first=True,
padding_value=vocab[""])
return padded, labels
内置数据集
from torchtext.datasets import IMDB, AG_NEWS, WikiText2
# IMDB(情感)
train_iter, test_iter = IMDB(split=("train", "test"))
# 每条: (label, text), label: 1=neg, 2=pos
# AG_NEWS(新闻分类,4 类)
train_iter, test_iter = AG_NEWS(split=("train", "test"))
# WikiText2(语言模型)
train_iter, valid_iter, test_iter = WikiText2(
split=("train", "valid", "test"))
# 迭代器 → 列表
train_data = list(train_iter)
print(len(train_data), train_data[0])
预训练词向量
from torchtext.vocab import GloVe
# 下载 GloVe 预训练向量
vectors = GloVe(name="6B", dim=100)
# name: 6B/42B/840B, dim: 50/100/200/300
# 查询
vectors["cat"] # 词向量(torch.Tensor)
vectors.stoi["cat"] # 词 → 索引
vectors.itos[10] # 索引 → 词
# 用预训练向量初始化 Embedding
def build_embedding(vocab, vectors, dim=100,
padding_idx=0):
import torch.nn as nn
import torch.nn.init as init
emb = nn.Embedding(len(vocab), dim,
padding_idx=padding_idx)
init.normal_(emb.weight, std=0.1)
# 匹配词表
for word, idx in vocab.get_stoi().items():
if word in vectors.stoi:
emb.weight.data[idx] = vectors[word]
emb.weight.data[padding_idx] = 0 # pad 归零
return emb
完整文本分类管线
# 组装: 分词 → 词表 → Dataset → DataLoader → 模型
class TextDataset(torch.utils.data.Dataset):
def __init__(self, data_iter, tokenizer,
vocab, max_len=200):
self.samples = []
for label, text in data_iter:
tokens = tokenizer(text)[:max_len]
ids = [vocab[t] for t in tokens]
ids += [vocab[""]] * (max_len - len(ids))
label = 1 if label == 2 else 0
self.samples.append(
(torch.tensor(ids), label))
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
return self.samples[idx]
train_ds = TextDataset(IMDB(split="train"),
tokenizer, vocab)
loader = torch.utils.data.DataLoader(
train_ds, batch_size=64, shuffle=True)
# 模型(Embedding + LSTM/Transformer)训练同前
torchtext 使用建议
# 1. 分词与词表: torchtext 提供标准方案
# 2. 简单任务: 自建 Dataset 更灵活
# 3. 复杂任务(预训练): 直接用 HuggingFace
# (tokenizer + 预训练模型一站式)
# 4. torchtext API 变动频繁,锁定版本
# 5. 大数据: 流式迭代器处理,不全部载入内存
torchtext 是 PyTorch 文本生态的入口:分词、词表、数据集标准化。掌握后,自定义文本任务的数据管线都能快速搭建。
PyTorch 混合精度训练
PyTorch 混合精度训练
混合精度训练(AMP)用 FP16 与 FP32 混合计算:显存减半、速度提升,精度基本无损,是现代训练的标准配置。
原理
# FP16 优点: 占一半显存,计算更快(GPU 原生加速)
# FP16 问题: 数值范围小,梯度易下溢为 0
# AMP 方案:
# - 前向/反向用 FP16(快、省显存)
# - 主权重保持 FP32(精度)
# - 损失缩放(loss scaling)防梯度下溢
# - 更新时转回 FP32
基础用法(推荐)
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# 开启 AMP
scaler = torch.cuda.amp.GradScaler() # 梯度缩放器
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
# 前向在 autocast 上下文(自动混合精度)
with torch.cuda.amp.autocast():
loss = criterion(model(x), y)
# 缩放损失 → 反传 → 缩放梯度更新
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
PyTorch 2.x 简化
# 2.0+ 更简洁
torch.set_autocast_enabled(True) # 全局开启
# 或上下文
with torch.autocast(device_type="cuda", dtype=torch.float16):
out = model(x)
loss = criterion(out, y)
# 2.3+ 推荐
torch.set_float32_matmul_precision("high")
# 自动混合精度训练也可用:
# torch.autocast 包裹整个训练循环
配合梯度裁剪
# 混合精度 + 梯度裁剪的坑:
# 缩放后的梯度范数 ≠ 真实范数
# 必须先 unscale 再裁剪
scaler.scale(loss).backward()
# 方式1: 让 scaler 处理裁剪
scaler.unscale_(optimizer) # 先还原梯度
torch.nn.utils.clip_grad_norm_(
model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
# 方式2: 传裁剪值给 step
scaler.step(optimizer,
clip_grad_norm=1.0) # 内部处理
检查损失是否溢出
# scaler 自动检测 inf/NaN 并跳过更新
# 可手动检查
if torch.isinf(loss) or torch.isnan(loss):
print("损失异常, 跳过该 batch")
# 继续下一个 batch 即可
# scaler.update() 会根据溢出情况
# 自动调整缩放因子(不缩放更新则放大)
CPU/非 CUDA 环境
# AMP 主要针对 GPU; CPU 用 bfloat16 有加速
with torch.autocast(device_type="cpu",
dtype=torch.bfloat16):
out = model(x)
# 判断可用性
torch.cuda.is_bf16_supported()
torch.cuda.amp.autocast(enabled=torch.cuda.is_available())
显存与速度收益
# 典型收益:
# - 显存: 降低 30-50%(可加大 batch)
# - 速度: 1.5-2.5x(取决于 GPU 的 FP16 支持)
# - 精度: 基本无损(大模型偶尔微降)
# 何时不用:
# - 数值敏感任务(个别场景)
# - CPU 训练(收益小)
# - 显式需要 FP32 精度
完整模板
import torch
def train_amp(model, loader, criterion,
optimizer, epochs, device):
model = model.to(device)
scaler = torch.cuda.amp.GradScaler()
for epoch in range(epochs):
model.train()
for x, y in loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
with torch.cuda.amp.autocast():
loss = criterion(model(x), y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Epoch {epoch+1}: {loss.item():.4f}")
# 推理时也可用 autocast 提速
model.eval()
with torch.cuda.amp.autocast(), torch.no_grad():
pred = model(x)
混合精度是 GPU 训练的"免费午餐":加三行代码,显存减半、速度翻倍。建议所有 CUDA 训练默认开启。
PyTorch TorchScript/ONNX 导出
PyTorch TorchScript/ONNX 导出
模型导出是部署的关键:把训练好的 PyTorch 模型转换为可跨平台推理的格式。本节介绍 TorchScript 与 ONNX 两种主流方式。
导出方式对比
TorchScript:
- PyTorch 原生格式(.pt)
- 支持 C++/Java/移动端
- tracing 或 scripting 两种方式
- 保真度高,复杂模型支持好
ONNX:
- 开放标准格式(.onnx)
- 可转到其他引擎(ONNX Runtime/TensorRT)
- 跨框架(ONNX → TF/TFLite/CoreML)
- 部署生态最广
选择:
- 纯 PyTorch 生态 → TorchScript
- 跨平台/异构部署 → ONNX
TorchScript 导出(Tracing)
import torch
model = MyModel().eval()
# 示例输入(决定输入形状)
example = torch.randn(1, 3, 224, 224)
# tracing: 跟踪一次前向,固化计算图
traced = torch.jit.trace(model, example)
traced.save("model_traced.pt")
# 加载与推理
loaded = torch.jit.load("model_traced.pt")
with torch.no_grad():
out = loaded(example)
# 适用: 无控制流(if/for 依赖数据)的模型
TorchScript 导出(Scripting)
# scripting: 直接编译模型源码
# 支持控制流,适合动态模型
scripted = torch.jit.script(model)
scripted.save("model_scripted.pt")
# 混合: 先用 script 打包动态部分
# 再 trace 外部(少用)
# 注意 scripting 要求:
# - 使用 torch 支持的语法
# - 类型标注(部分情况)
# - 避免 Python 特有特性(如 lambda)
ONNX 导出
model = MyModel().eval()
example = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model,
example,
"model.onnx",
export_params=True, # 导出参数
opset_version=17, # ONNX 算子集版本
do_constant_folding=True, # 常量折叠优化
input_names=["input"],
output_names=["output"],
dynamic_axes={ # 动态维度
"input": {0: "batch_size"},
"output": {0: "batch_size"},
},
)
# 验证导出
import onnx
onnx.checker.check_model(onnx_model)
onnx.helper.printable_graph(onnx_model.graph)
ONNX Runtime 推理
import onnxruntime as ort
import numpy as np
# 创建会话(可指定 providers)
sess = ort.InferenceSession(
"model.onnx",
providers=["CUDAExecutionProvider",
"CPUExecutionProvider"])
# 输入格式
input_name = sess.get_inputs()[0].name
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 推理
outputs = sess.run(None, {input_name: x})
pred = outputs[0]
print(pred.shape) # [1, num_classes]
TensorRT 优化(可选)
# 高性能推理: ONNX → TensorRT
# 方式1: trtexec 命令行
trtexec --onnx=model.onnx --saveEngine=model.trt --fp16
# 方式2: Python API
import tensorrt as trt
# 构建引擎(略)
# TensorRT 在 NVIDIA GPU 上延迟最低
# 其他引擎:
# OpenVINO(Intel)
# CoreML(Apple)
# TFLite(移动端)
导出常见问题
# 1. 动态形状: 必须声明 dynamic_axes
# 2. 不支持算子: 升级 opset/简化模型
# 或 torch.onnx.export 前 onnxsim 简化
# python -m onnxsim model.onnx model_sim.onnx
# 3. 预处理一致性: 导出后的输入要求与训练相同
# (归一化在模型外或模型内,保持一致)
# 4. 验证: 导出前后输出对比(误差 < 1e-3)
def verify(model, onnx_model, x):
with torch.no_grad():
torch_out = model(x).numpy()
ort_out = onnx_session.run(None, {"input": x.numpy()})[0]
diff = np.abs(torch_out - ort_out).max()
print(f"最大差异: {diff:.2e}") # 应很小
导出最佳实践
# 1. model.eval() + no_grad 后导出
# 2. 输入形状明确: 固定或声明动态轴
# 3. 导出前验证输出一致
# 4. 版本记录: 模型+代码+依赖版本
# 5. 量化导出(部署优化):
# - 动态量化(CPU): torch.quantization
# - INT8 量化(TensorRT/ORT)
# 模型版本管理:
# 导出时附 config.json(结构/预处理/版本)
TorchScript 保 PyTorch 生态,ONNX 通全平台。导出-验证-部署三步走,模型才能真正上线服务。
PyTorch 分布式训练
PyTorch 分布式训练
单卡不够用时,分布式训练扩展规模:多 GPU、多机。本节介绍数据并行(DP/DDP)与相关概念。
并行策略
# 1. 数据并行(Data Parallel)
# 每卡一份模型副本,分不同 batch 数据
# 梯度同步更新 → 扩大 batch,加速
# 最常用,简单有效
# 2. 模型并行(Model Parallel)
# 模型分片到多卡(超大模型)
# 层/张量切分,单卡放不下时用
# 3. 流水线并行(Pipeline)
# 层分阶段,数据流式经过各卡
# 4. FSDP(全分片数据并行)
# 参数/梯度/优化器分片,大模型首选
# (HuggingFace 大模型训练标配)
DataParallel(单机多卡,简单)
import torch.nn as nn
model = MyModel()
model = nn.DataParallel(model) # 自动分配到多卡
model = model.to("cuda")
# 训练代码不变,自动并行
# 注意: 实际 batch 被分成 n 份
# 例如 4 卡, batch=64 → 每卡 16
# 缺点:
# - 主卡负载高(梯度同步瓶颈)
# - 扩展性一般(单机)
# 小规模可用,推荐 DDP
DistributedDataParallel(推荐)
# 文件方式训练脚本 train_ddp.py
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
def train(rank, world_size):
# 1. 初始化进程组
dist.init_process_group(
backend="nccl", # GPU 用 nccl
init_method="env://",
rank=rank,
world_size=world_size,
)
torch.cuda.set_device(rank)
# 2. 模型
model = MyModel().to(rank)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[rank])
# 3. 数据(分布式采样器!)
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=64,
sampler=sampler,
num_workers=4)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
sampler.set_epoch(epoch) # 每轮打乱
for x, y in loader:
x, y = x.to(rank), y.to(rank)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
dist.destroy_process_group()
if __name__ == "__main__":
world_size = torch.cuda.device_count()
mp.spawn(train, args=(world_size,),
nprocs=world_size)
启动方式
# 方式1: torchrun(推荐)
torchrun --nproc_per_node=4 train_ddp.py
# 方式2: python -m torch.distributed.launch
python -m torch.distributed.launch --nproc_per_node=4 train_ddp.py
# 方式3: 多机(需共享文件系统)
# 每台机器
torchrun --nnodes=2 --nproc_per_node=8 --rdzv_endpoint=主节点IP:29500 --rdzv_backend=c10d train_ddp.py
# SLURM 集群
# srun python train_ddp.py
DDP 关键点
# 1. 每进程一个 GPU(rank 对应 device)
# 2. DistributedSampler 保证数据不重复
# 3. sampler.set_epoch(epoch) 每轮打乱
# 4. 模型保存: 只在一个 rank 保存
if rank == 0:
torch.save(model.module.state_dict(), "model.pth")
# 5. 日志打印: 只让 rank 0 输出
if rank == 0:
print(...)
# 6. batch_size 是"每卡"大小
# 总 batch = batch_size * world_size
# 注意调 lr(linear scaling 等)
梯度与批大小策略
# 同步 BatchNorm(跨卡统计)
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
# 学习率缩放(线性缩放规则):
# 新 lr = 原 lr * (新总batch / 原总batch)
# 混合精度 + DDP(兼容)
scaler = torch.cuda.amp.GradScaler()
# 训练循环同 AMP 章节,注意每个进程独立 scaler
FSDP(大模型训练)
# 全分片数据并行: 参数/梯度分片到各卡
# 适合 1B+ 参数模型
from torch.distributed.fsdp import (
FullyShardedDataParallel as FSDP,
CPUOffload,
)
from torch.distributed.fsdp.wrap import (
transformer_auto_wrap_policy)
model = FSDP(
model,
auto_wrap_policy=transformer_auto_wrap_policy(
transformer_layer_cls={TransformerBlock}),
cpu_offload=CPUOffload(offload_params=True),
)
# 或直接用 HuggingFace accelerate:
# from accelerate import Accelerator
# accelerator = Accelerator()
# model, opt, loader = accelerator.prepare(...)
分布式训练最佳实践
# 1. 先单卡跑通,再上多卡
# 2. 小规模用 DataParallel,大规模用 DDP
# 3. 数据用 DistributedSampler(不要手动切)
# 4. 保存/日志只在 rank 0
# 5. 固定种子保证可复现
# 6. 检查 GPU 利用率(nvidia-smi)
# 7. 大模型用 FSDP/accelerate
# 监控
nvidia-smi # GPU 状态
nvidia-smi dmon # 实时监控
# 利用率低 → 检查 num_workers / 数据管线
分布式训练让"算力"可扩展:DP 入门、DDP 标准、FSDP 扛大模型。先会用 torchrun 跑 DDP,就能扩展到多机与超大模型。
PyTorch 注意力机制
PyTorch 注意力机制
注意力(Attention)是现代深度学习的核心思想,也是 Transformer 的基石。本节深入实现各类注意力,理解其本质。
注意力本质
# 核心问题: 处理序列时,如何"关注"重要的部分?
# 公式: Attention(Q,K,V) = softmax(QK^T/√d) V
# Q(查询): 我想找什么
# K(键): 我有的是什么(匹配查询)
# V(值): 真正要取的内容
# 流程:
# 1. Q 与所有 K 算相似度
# 2. softmax 归一化为权重
# 3. 权重加权求和 V → 输出
# 直觉: 每个位置从全序列"挑重点"聚合信息
缩放点积注意力
import torch
import torch.nn as nn
import torch.nn.functional as F
def attention(q, k, v, mask=None, dropout=None):
"""缩放点积注意力
q/k/v: [batch, heads, seq, dim]
"""
d_k = q.size(-1)
scores = q @ k.transpose(-2, -1) / (d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
weights = F.softmax(scores, dim=-1)
if dropout:
weights = dropout(weights)
return weights @ v, weights
多头注意力
class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, n_heads=8,
dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.wq = nn.Linear(d_model, d_model)
self.wk = nn.Linear(d_model, d_model)
self.wv = nn.Linear(d_model, d_model)
self.wo = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
batch, seq, _ = x.shape
# 线性变换并分头
q = self.wq(x).view(batch, seq,
self.n_heads, self.d_k)
q = q.transpose(1, 2) # [b, h, s, d]
k = self.wk(x).view(batch, seq,
self.n_heads, self.d_k).transpose(1, 2)
v = self.wv(x).view(batch, seq,
self.n_heads, self.d_k).transpose(1, 2)
# 注意力
attn, _ = attention(q, k, v, mask, self.dropout)
# 合并头
attn = attn.transpose(1, 2).contiguous()
attn = attn.view(batch, seq, self.d_model)
return self.wo(attn)
# 多头意义: 不同头关注不同关系(语法/语义/位置)
# 例: 一个头关注相邻词,一个头关注远距离宾语
自注意力与交叉注意力
# 自注意力(Self-Attention)
# Q=K=V(都来自同一序列)
# 用途: 编码器(序列内部建模)
self_attn = MultiHeadAttention(d_model=512, n_heads=8)
out = self_attn(x) # x: [b, s, 512]
# 交叉注意力(Cross-Attention)
# Q 来自解码器, K/V 来自编码器输出
# 用途: 解码器"查询"编码器的信息
cross_attn = MultiHeadAttention(d_model=512, n_heads=8)
out = cross_attn(q=dec_out, k=enc_out, v=enc_out)
# 注意: 上述实现需支持 q/k/v 不同输入
掩码注意力
# 因果掩码(生成任务: 只看过去)
def causal_mask(seq_len):
"""下三角掩码,未来位置为 False"""
return torch.tril(torch.ones(seq_len, seq_len)
).bool()
# 填充掩码(忽略 padding)
def padding_mask(tokens, pad_idx=0):
"""padding 位置为 False"""
return (tokens != pad_idx).unsqueeze(1).unsqueeze(2)
# 组合(解码器用):
# 因果掩码 + 填充掩码(AND)
seq = 20
mask = causal_mask(seq)
# 使用时: mask 形状需匹配 [b*h, s, s] 或广播
PyTorch 内置实现
# 用 F.scaled_dot_product_attention(2.0+,推荐)
out = F.scaled_dot_product_attention(
q, k, v, # [b, h, s, d]
attn_mask=mask,
dropout_p=0.1,
is_causal=True, # 因果掩码快捷参数
)
# 自动选择高效内核(Flash Attention/Memory-Efficient)
# 或 nn.MultiheadAttention(完整封装)
mha = nn.MultiheadAttention(
embed_dim=512, num_heads=8, batch_first=True)
out, attn_weights = mha(x, x, x)
注意力可视化
# 查看注意力权重(解释模型)
import matplotlib.pyplot as plt
_, attn_weights = mha(q, k, v)
# attn_weights: [batch, heads, seq_q, seq_k]
def plot_attention(weights, tokens, head=0):
w = weights[0, head].detach().numpy() # [s, s]
plt.imshow(w, cmap="viridis")
plt.xticks(range(len(tokens)), tokens, rotation=90)
plt.yticks(range(len(tokens)), tokens)
plt.colorbar()
plt.show()
# tokens = ["我", "爱", "学习", "编程"]
# plot_attention(attn_weights, tokens)
注意力变体
# 1. 多头注意力(MHA): Transformer 标准
# 2. 分组查询注意力(GQA): 共享 KV,省显存
# (Llama 2/3 使用)
# 3. 多查询注意力(MQA): 单 KV,更省
# 4. Flash Attention: IO 优化,快速省显存
# 5. 稀疏/线性注意力: 长序列降复杂度
# 6. 交叉注意力: 跨序列交互
# 长序列(Swin/长文本):
# 窗口注意力/滑动窗口(降低复杂度)
# 现代实践: 直接用 F.scaled_dot_product_attention
# 自动选最优内核,无需手写细节
注意力机制是理解现代 AI(GPT/BERT)的钥匙。掌握缩放点积、多头与掩码三个概念,Transformer 就不再神秘。