PyTorch深度学习全栈指南

PyTorch 完整知识总结与使用教程

---

目录

1. [PyTorch 简介与安装](#PyTorch 简介与安装)
2. 核心数据结构：张量（Tensor）
3. 自动微分（Autograd）
4. 神经网络构建（torch.nn）
5. 数据加载与处理
6. 损失函数与优化器
7. 完整训练工作流
8. [GPU 加速](#GPU 加速)
9. 模型保存与加载
10. 迁移学习
11. 分布式训练
12. [模型部署（TorchScript & ONNX）](#模型部署（TorchScript & ONNX）)
13. [PyTorch 2.x 新特性：torch.compile](#PyTorch 2.x 新特性：torch.compile)
14. 常用工具与生态
15. 最佳实践与调试技巧

---

1. PyTorch 简介与安装

1.1 什么是 PyTorch？

PyTorch 是由 Meta AI （Facebook）研发、现由 Linux 基金会 托管的开源深度学习框架。它基于 Lua 版本的 Torch 库，使用 Python 和 C++ 编写，是当今学术界和工业界最主流的深度学习框架之一。

核心两大特性：

┌──────────────────────────────────────────────────────────┐
│                    PyTorch 核心特性                       │
├─────────────────────────┬────────────────────────────────┤
│  动态计算图（Define-by-Run）│  自动微分（Autograd）          │
│  运行时构建，灵活调试      │  自动计算梯度，支持反向传播       │
├─────────────────────────┴────────────────────────────────┤
│  GPU 加速（CUDA）  ｜  与 NumPy 无缝集成  ｜ Pythonic API │
└──────────────────────────────────────────────────────────┘

PyTorch vs TensorFlow 核心区别：

特性	PyTorch	TensorFlow
计算图	动态图（运行时构建）	静态图（v1）/ 动态图（v2 Eager）
调试方式	原生 Python 调试器	专用调试工具
学习曲线	较平缓，Pythonic	较陡峭
研究领域	首选框架	生产环境历史优势
部署	TorchScript/ONNX	TensorFlow Serving/TFLite

1.2 安装

官方推荐安装方式（访问 pytorch.org 获取最新命令）：

# CPU 版本
pip install torch torchvision torchaudio

# CUDA 12.1 GPU 版本（以 CUDA 12.1 为例）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Conda 安装
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

验证安装：

import torch

print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 是否可用: {torch.cuda.is_available()}")
print(f"CUDA 版本: {torch.version.cuda}")
print(f"GPU 设备数量: {torch.cuda.device_count()}")

---

2. 核心数据结构：张量（Tensor）

2.1 张量的概念

张量（Tensor）是 PyTorch 中最基本的数据结构，可以理解为 N 维数组 ，是 NumPy ndarray 的 GPU 增强版。

维度关系示意图：
  0维 (标量)：  42
  1维 (向量)：  [1, 2, 3]
  2维 (矩阵)：  [[1, 2], [3, 4]]
  3维 (RGB图像)：shape = (C, H, W) = (3, 224, 224)
  4维 (批量图像)：shape = (N, C, H, W) = (32, 3, 224, 224)

2.2 创建张量

import torch
import numpy as np

# ── 基础创建方式 ──────────────────────────────────────────
# 从 Python 列表创建
t1 = torch.tensor([[1, 2, 3], [4, 5, 6]])

# 从 NumPy 创建（共享内存！）
np_arr = np.array([1.0, 2.0, 3.0])
t2 = torch.from_numpy(np_arr)

# ── 特殊张量 ──────────────────────────────────────────────
t_zeros = torch.zeros(3, 4)          # 全零张量
t_ones  = torch.ones(3, 4)           # 全一张量
t_rand  = torch.rand(3, 4)           # 均匀分布 [0, 1)
t_randn = torch.randn(3, 4)          # 标准正态分布
t_eye   = torch.eye(4)               # 单位矩阵
t_arange = torch.arange(0, 10, 2)   # [0, 2, 4, 6, 8]
t_linspace = torch.linspace(0, 1, 5) # 等差数列

# ── 和已有张量形状相同的张量 ────────────────────────────────
t_like_zeros = torch.zeros_like(t1, dtype=torch.float32)
t_like_rand  = torch.rand_like(t_like_zeros)

2.3 张量属性

t = torch.randn(3, 4, dtype=torch.float32)

print(f"形状 (shape):   {t.shape}")         # torch.Size([3, 4])
print(f"数据类型(dtype): {t.dtype}")         # torch.float32
print(f"存储设备(device):{t.device}")        # cpu
print(f"维度数量(ndim):  {t.ndim}")          # 2
print(f"元素总数(numel): {t.numel()}")       # 12
print(f"需要梯度:        {t.requires_grad}") # False

常用数据类型：

类型	说明	别名
torch.float32	32位浮点（默认）	torch.float
torch.float64	64位浮点	torch.double
torch.float16	16位浮点（半精度）	torch.half
torch.bfloat16	Brain 浮点格式，训练常用	---
torch.int32	32位整数	torch.int
torch.int64	64位整数	torch.long
torch.bool	布尔类型	---

2.4 张量操作

# ── 形状变换 ───────────────────────────────────────────────
t = torch.randn(4, 4)

t_view    = t.view(16)           # 共享内存，连续张量
t_reshape = t.reshape(2, 8)      # 不保证共享内存
t_flat    = t.flatten()          # 完全展平
t_unsq    = t.unsqueeze(0)       # 在第0维增加一个维度 → [1, 4, 4]
t_sq      = t_unsq.squeeze(0)    # 删除大小为1的维度 → [4, 4]
t_T       = t.T                  # 转置（2D）
t_perm    = t.permute(1, 0)      # 维度置换（多维通用）

# ── 数学运算 ───────────────────────────────────────────────
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

# 元素级运算
c = a + b               # 等同 torch.add(a, b)
c = a * b               # 元素级乘法
c = a / b               # 元素级除法
c = a ** 2              # 幂运算

# 矩阵运算
A = torch.randn(3, 4)
B = torch.randn(4, 5)
C = A @ B               # 矩阵乘法，等同 torch.matmul(A, B)
C = torch.mm(A, B)      # 2D 矩阵乘法

# 聚合运算
t.sum()           # 所有元素求和
t.sum(dim=0)      # 沿第0维求和
t.mean()          # 均值
t.max()           # 最大值
t.argmax(dim=1)   # 最大值的索引
t.std()           # 标准差

# ── 索引与切片（类似 NumPy） ────────────────────────────────
t = torch.arange(12).reshape(3, 4)
print(t[0])          # 第0行
print(t[:, 1])       # 第1列
print(t[1:, 2:])     # 子矩阵
print(t[t > 5])      # 布尔索引

# ── 拼接与堆叠 ─────────────────────────────────────────────
t1 = torch.zeros(2, 3)
t2 = torch.ones(2, 3)

cat_0 = torch.cat([t1, t2], dim=0)   # 沿第0维拼接 → [4, 3]
cat_1 = torch.cat([t1, t2], dim=1)   # 沿第1维拼接 → [2, 6]
stack = torch.stack([t1, t2], dim=0) # 新增维度堆叠 → [2, 2, 3]

# ── 与 NumPy 互转 ──────────────────────────────────────────
t = torch.ones(5)
n = t.numpy()          # Tensor → NumPy（CPU张量，共享内存）
t2 = torch.from_numpy(n)  # NumPy → Tensor

2.5 张量的内存机制

# view() 共享内存 ------ 改一个，另一个也变！
a = torch.randn(4, 4)
b = a.view(16)
b[0] = 999
print(a[0, 0])  # 999（已改变）

# 推荐写法：需要独立副本时先 clone()
b = a.clone().view(16)
b[0] = 999
print(a[0, 0])  # 原值不变

# contiguous() ------ 某些操作后张量可能不连续
t_transposed = a.T
t_contiguous = t_transposed.contiguous()  # 使内存连续

---

3. 自动微分（Autograd）

3.1 计算图与梯度

PyTorch 使用 动态计算图（DAG） 记录前向传播中的所有操作，以支持反向传播自动计算梯度。

前向传播示意图：

  x ─→ [操作1] ─→ [操作2] ─→ ... ─→ loss
  ↑        ↑          ↑
  叶节点   中间节点     根节点

backward() 从 loss 出发，沿箭头反方向传递梯度（链式法则）

3.2 requires_grad 与梯度计算

import torch

# 创建需要追踪梯度的张量（叶节点）
x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)
b = torch.tensor([1.0], requires_grad=True)

# 前向传播
y = w * x + b      # y = 3*2 + 1 = 7
loss = y ** 2      # loss = 49

# 反向传播
loss.backward()

# 查看梯度
# d(loss)/dx = 2*y * w = 2*7*3 = 42
print(f"x.grad = {x.grad}")  # tensor([42.])
print(f"w.grad = {w.grad}")  # tensor([28.])  = 2*7*2
print(f"b.grad = {b.grad}")  # tensor([14.])  = 2*7*1

3.3 梯度控制

# ── 方式1：with torch.no_grad()（推理时常用）────────────────
model.eval()
with torch.no_grad():
    output = model(input)   # 不追踪梯度，节省内存和计算

# ── 方式2：@torch.no_grad() 装饰器 ────────────────────────
@torch.no_grad()
def predict(model, x):
    return model(x)

# ── 方式3：detach()（从计算图中分离）──────────────────────
x = torch.randn(3, requires_grad=True)
y = x * 2
z = y.detach()   # z 不再追踪梯度
print(z.requires_grad)  # False

# ── 梯度清零（每次 backward 前必做！）──────────────────────
# 注意：PyTorch 梯度是累加的，不清零会叠加！
optimizer.zero_grad()   # 通过优化器清零（推荐）
# 或
for param in model.parameters():
    if param.grad is not None:
        param.grad.zero_()

3.4 计算图特性

# DAG 是动态的------每次 forward 重新构建
# 这意味着可以在 forward 中使用 Python 控制流！
import torch

def dynamic_model(x, use_relu=True):
    if use_relu:           # 动态控制流
        return torch.relu(x)
    else:
        return torch.sigmoid(x)

x = torch.randn(5, requires_grad=True)
y = dynamic_model(x, use_relu=True)
y.sum().backward()

---

4. 神经网络构建（torch.nn）

4.1 nn.Module 基类

所有神经网络模型都继承自 torch.nn.Module，只需实现：

• __init__：定义网络层
• forward：定义前向传播逻辑（backward 由 autograd 自动实现）

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNet, self).__init__()
        # 定义网络层
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        # 定义数据流向
        x = self.fc1(x)
        x = self.bn1(x)
        x = F.relu(x)           # 激活函数
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 实例化
model = SimpleNet(784, 256, 10)
print(model)

# 查看参数数量
total_params = sum(p.numel() for p in model.parameters())
print(f"总参数量: {total_params:,}")

4.2 常用层（Layers）

import torch.nn as nn

# ── 线性层（全连接层）─────────────────────────────────────
nn.Linear(in_features=128, out_features=64, bias=True)

# ── 卷积层 ─────────────────────────────────────────────────
nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3)
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, 
          stride=1, padding=1)   # 图像卷积
nn.Conv3d(...)                   # 视频/体积数据

# 反卷积（上采样）
nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=4, stride=2)

# ── 池化层 ─────────────────────────────────────────────────
nn.MaxPool2d(kernel_size=2, stride=2)
nn.AvgPool2d(kernel_size=2)
nn.AdaptiveAvgPool2d(output_size=(1, 1))  # 全局平均池化

# ── 归一化层 ────────────────────────────────────────────────
nn.BatchNorm1d(num_features=64)   # 1D 批归一化
nn.BatchNorm2d(num_features=64)   # 2D 批归一化（图像）
nn.LayerNorm(normalized_shape=64) # 层归一化（NLP常用）
nn.GroupNorm(num_groups=8, num_channels=64)

# ── Dropout ─────────────────────────────────────────────────
nn.Dropout(p=0.5)    # 1D
nn.Dropout2d(p=0.5)  # 2D（通道级别）

# ── 循环层 ─────────────────────────────────────────────────
nn.RNN(input_size=10, hidden_size=20, num_layers=2)
nn.LSTM(input_size=10, hidden_size=20, num_layers=2, 
        batch_first=True, bidirectional=True)
nn.GRU(input_size=10, hidden_size=20)

# ── Embedding ───────────────────────────────────────────────
nn.Embedding(num_embeddings=10000, embedding_dim=128)

# ── Transformer ─────────────────────────────────────────────
nn.MultiheadAttention(embed_dim=512, num_heads=8)
nn.TransformerEncoder(...)
nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6)

4.3 激活函数

import torch.nn.functional as F

# 常用激活函数
F.relu(x)          # ReLU: max(0, x)
F.sigmoid(x)       # Sigmoid: 1/(1+e^-x)，输出(0,1)
F.tanh(x)          # Tanh: (-1,1)
F.softmax(x, dim=1)  # Softmax（多分类输出层）
F.log_softmax(x, dim=1)  # LogSoftmax
F.gelu(x)          # GELU（Transformer常用）
F.leaky_relu(x, negative_slope=0.01)  # Leaky ReLU
F.elu(x)           # ELU
F.selu(x)          # SELU

# 作为 nn.Module 使用
nn.ReLU(inplace=True)   # inplace=True 节省内存
nn.GELU()
nn.Sigmoid()
nn.Tanh()
nn.Softmax(dim=1)

4.4 使用 nn.Sequential 快速搭建

import torch.nn as nn

# 方式一：直接传入列表
model = nn.Sequential(
    nn.Linear(784, 512),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# 方式二：使用 OrderedDict 命名层
from collections import OrderedDict

model = nn.Sequential(OrderedDict([
    ('fc1',    nn.Linear(784, 512)),
    ('relu1',  nn.ReLU()),
    ('drop1',  nn.Dropout(0.3)),
    ('fc2',    nn.Linear(512, 10))
]))

# 访问特定层
print(model.fc1)
print(model[0])      # 按索引访问

4.5 经典网络结构示例：CNN

import torch.nn as nn
import torch.nn.functional as F

class ConvNet(nn.Module):
    """简单 CNN 用于图像分类（以 CIFAR-10 为例）"""
    
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        # 特征提取器
        self.features = nn.Sequential(
            # Block 1: [3, 32, 32] → [32, 32, 32]
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),   # → [32, 16, 16]
            
            # Block 2: [32, 16, 16] → [64, 8, 8]
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),   # → [64, 8, 8]
        )
        # 分类器
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((4, 4)),  # 全局自适应池化
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

model = ConvNet(num_classes=10)

---

5. 数据加载与处理

5.1 Dataset 与 DataLoader 的关系

数据管道架构：

  原始数据 (文件/数据库)
       ↓
  Dataset（定义如何读取单个样本）
       ↓
  DataLoader（批量加载 + 多进程 + shuffle）
       ↓
  训练循环

5.2 自定义 Dataset

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os

class ImageDataset(Dataset):
    """自定义图像数据集"""
    
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        # 获取所有图片路径和标签
        self.samples = self._load_samples()
    
    def _load_samples(self):
        """收集所有样本路径和标签"""
        samples = []
        classes = sorted(os.listdir(self.img_dir))  # 类别目录
        self.class_to_idx = {c: i for i, c in enumerate(classes)}
        
        for cls in classes:
            cls_dir = os.path.join(self.img_dir, cls)
            for fname in os.listdir(cls_dir):
                if fname.endswith(('.jpg', '.png')):
                    samples.append((
                        os.path.join(cls_dir, fname),
                        self.class_to_idx[cls]
                    ))
        return samples
    
    def __len__(self):
        """返回数据集大小（必须实现）"""
        return len(self.samples)
    
    def __getitem__(self, idx):
        """返回单个样本（必须实现）"""
        img_path, label = self.samples[idx]
        image = Image.open(img_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

5.3 数据预处理与增强（transforms）

import torchvision.transforms as transforms

# 训练集变换（包含数据增强）
train_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomCrop(224),               # 随机裁剪
    transforms.RandomHorizontalFlip(p=0.5),   # 随机水平翻转
    transforms.RandomRotation(degrees=15),    # 随机旋转
    transforms.ColorJitter(
        brightness=0.2, contrast=0.2,
        saturation=0.2, hue=0.1
    ),
    transforms.ToTensor(),                    # PIL → Tensor [0,1]
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],           # ImageNet 均值
        std=[0.229, 0.224, 0.225]             # ImageNet 标准差
    )
])

# 验证/测试集变换（不做增强）
val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

5.4 DataLoader 配置

from torch.utils.data import DataLoader, random_split

# 使用内置数据集（以 FashionMNIST 为例）
from torchvision import datasets

full_dataset = datasets.FashionMNIST(
    root='./data',
    train=True,
    download=True,
    transform=train_transform
)

# 分割训练集和验证集
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])

# 创建 DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=64,          # 每批次样本数
    shuffle=True,           # 训练时打乱
    num_workers=4,          # 多进程加载（Linux/Mac 推荐 4-8）
    pin_memory=True,        # 锁页内存，加速 GPU 传输
    drop_last=True          # 丢弃最后不完整的批次
)

val_loader = DataLoader(
    val_dataset,
    batch_size=128,
    shuffle=False,          # 验证时不需要打乱
    num_workers=4,
    pin_memory=True
)

# 遍历 DataLoader
for batch_idx, (images, labels) in enumerate(train_loader):
    print(f"批次 {batch_idx}: images={images.shape}, labels={labels.shape}")
    break  # 只看第一批

5.5 内置数据集

from torchvision import datasets

# 常用视觉数据集
datasets.MNIST(root, train, transform, download)
datasets.FashionMNIST(root, train, transform, download)
datasets.CIFAR10(root, train, transform, download)
datasets.CIFAR100(root, train, transform, download)
datasets.ImageNet(root, split='train', transform=transform)
datasets.ImageFolder(root, transform)   # 自定义目录结构

# 文本数据集（torchtext）
# 音频数据集（torchaudio）

---

6. 损失函数与优化器

6.1 常用损失函数

import torch.nn as nn

# ── 分类任务 ───────────────────────────────────────────────
# 交叉熵损失（最常用，内部含 Softmax）
criterion = nn.CrossEntropyLoss()
# output: [N, C]，target: [N]（类别索引，非 one-hot）
loss = criterion(output, target)

# 二元交叉熵（二分类）
criterion = nn.BCELoss()          # 输入需先经过 Sigmoid
criterion = nn.BCEWithLogitsLoss()  # 更稳定，内部含 Sigmoid（推荐）

# 负对数似然（配合 LogSoftmax 使用）
criterion = nn.NLLLoss()

# ── 回归任务 ───────────────────────────────────────────────
criterion = nn.MSELoss()          # 均方误差
criterion = nn.L1Loss()           # 平均绝对误差（MAE）
criterion = nn.SmoothL1Loss()     # Huber Loss（对异常值鲁棒）

# ── 其他 ────────────────────────────────────────────────────
criterion = nn.KLDivLoss()        # KL 散度（知识蒸馏）
criterion = nn.TripletMarginLoss()  # 三元组损失（度量学习）
criterion = nn.CTCLoss()          # 序列识别（语音/OCR）

损失函数选择指南：

任务类型          推荐损失函数
────────────────────────────────────────────
二分类           BCEWithLogitsLoss
多分类           CrossEntropyLoss
多标签分类       BCEWithLogitsLoss
回归             MSELoss / SmoothL1Loss
序列到序列       CTCLoss
度量学习         TripletMarginLoss / ContrastiveLoss

6.2 优化器

import torch.optim as optim

# ── 常用优化器 ─────────────────────────────────────────────
# SGD（随机梯度下降，支持动量和权重衰减）
optimizer = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    weight_decay=1e-4
)

# Adam（自适应学习率，最常用）
optimizer = optim.Adam(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0
)

# AdamW（Adam + 解耦权重衰减，Transformer 首选）
optimizer = optim.AdamW(
    model.parameters(),
    lr=1e-4,
    weight_decay=0.01
)

# RMSprop
optimizer = optim.RMSprop(model.parameters(), lr=0.001)

# ── 为不同层设置不同学习率 ───────────────────────────────────
optimizer = optim.Adam([
    {'params': model.backbone.parameters(), 'lr': 1e-4},  # 主干网络
    {'params': model.head.parameters(),     'lr': 1e-3},  # 分类头
], lr=1e-3)

6.3 学习率调度器

import torch.optim.lr_scheduler as lr_scheduler

# StepLR：每 step_size 个 epoch，lr *= gamma
scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

# MultiStepLR：在指定 milestones 处衰减
scheduler = lr_scheduler.MultiStepLR(
    optimizer, milestones=[30, 60, 80], gamma=0.1
)

# CosineAnnealingLR：余弦退火（推荐）
scheduler = lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100, eta_min=1e-6
)

# ReduceLROnPlateau：监控指标，不改善就降低 lr
scheduler = lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=5
)

# OneCycleLR：超收敛（FastAI 方法）
scheduler = lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.01,
    steps_per_epoch=len(train_loader), epochs=50
)

# Warmup + Cosine（现代训练常用）
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
warmup = LinearLR(optimizer, start_factor=0.1, total_iters=5)
cosine = CosineAnnealingLR(optimizer, T_max=45)
scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[5])

---

7. 完整训练工作流

7.1 标准训练循环

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

def train_one_epoch(model, loader, criterion, optimizer, device):
    """训练一个 epoch"""
    model.train()  # 切换到训练模式（启用 BatchNorm / Dropout）
    
    total_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (inputs, targets) in enumerate(loader):
        # 1. 数据移到设备
        inputs  = inputs.to(device)
        targets = targets.to(device)
        
        # 2. 清零梯度（每批次必做！）
        optimizer.zero_grad()
        
        # 3. 前向传播
        outputs = model(inputs)
        loss    = criterion(outputs, targets)
        
        # 4. 反向传播
        loss.backward()
        
        # 5. 梯度裁剪（防止梯度爆炸，RNN/Transformer 常用）
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        # 6. 更新参数
        optimizer.step()
        
        # 统计指标
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total   += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    
    avg_loss = total_loss / len(loader)
    accuracy = 100. * correct / total
    return avg_loss, accuracy


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    """在验证/测试集上评估"""
    model.eval()  # 切换到评估模式（禁用 BatchNorm / Dropout）
    
    total_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in loader:
        inputs  = inputs.to(device)
        targets = targets.to(device)
        
        outputs = model(inputs)
        loss    = criterion(outputs, targets)
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total   += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    
    avg_loss = total_loss / len(loader)
    accuracy = 100. * correct / total
    return avg_loss, accuracy


# ── 完整训练流程 ─────────────────────────────────────────────
def train(config):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 初始化
    model     = ConvNet(num_classes=10).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs'])
    
    best_acc = 0.0
    
    for epoch in range(config['epochs']):
        # 训练
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )
        # 验证
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
        # 更新学习率
        scheduler.step()
        
        print(f"Epoch [{epoch+1:3d}/{config['epochs']}] "
              f"Train Loss: {train_loss:.4f} Acc: {train_acc:.2f}% | "
              f"Val Loss: {val_loss:.4f} Acc: {val_acc:.2f}% | "
              f"LR: {scheduler.get_last_lr()[0]:.6f}")
        
        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')
            print(f"  ✓ 保存最佳模型（Val Acc: {best_acc:.2f}%）")
    
    print(f"\n训练完成！最佳验证精度: {best_acc:.2f}%")

config = {'epochs': 50}

7.2 混合精度训练（AMP）

使用 FP16/BF16 混合精度，可将训练速度提升 2-4x，显存减少约一半：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # 梯度缩放器，防止 FP16 下溢

for inputs, targets in train_loader:
    inputs, targets = inputs.to(device), targets.to(device)
    
    optimizer.zero_grad()
    
    # 在 autocast 上下文中自动选择精度
    with autocast(dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    
    # 反向传播（使用缩放后的梯度）
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

---

8. GPU 加速

8.1 设备管理

import torch

# 自动选择设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 指定 GPU
device = torch.device('cuda:0')  # 第一块 GPU
device = torch.device('cuda:1')  # 第二块 GPU

# Apple Silicon
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

# 查询 GPU 信息
if torch.cuda.is_available():
    print(f"GPU 数量: {torch.cuda.device_count()}")
    print(f"当前 GPU: {torch.cuda.get_device_name(0)}")
    print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    print(f"显存已用: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
    print(f"显存缓存: {torch.cuda.memory_reserved(0) / 1e9:.2f} GB")

8.2 数据和模型的 GPU 迁移

# 模型迁移到 GPU
model = model.to(device)
# 或
model = model.cuda()

# 张量迁移
x = torch.randn(10, 3, 224, 224).to(device)
x = x.cuda()  # 等价写法

# 注意：不同设备的张量不能直接运算！
# 错误示例：
# cpu_tensor = torch.ones(5)
# gpu_tensor = torch.ones(5).cuda()
# result = cpu_tensor + gpu_tensor  # 报错！

# 正确做法：确保在同一设备
cpu_tensor = torch.ones(5)
gpu_tensor = cpu_tensor.to(device)   # 移到 GPU
result = gpu_tensor + gpu_tensor      # ✓

8.3 多 GPU 训练

# DataParallel（简单但效率较低）
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)  # 自动分发到所有可用 GPU
model = model.to(device)

# DistributedDataParallel（推荐，性能更好）
# 见第 11 节分布式训练

---

9. 模型保存与加载

9.1 保存和加载模型参数（推荐）

# ── 保存 ─────────────────────────────────────────────────────
# 只保存参数字典（推荐方式）
torch.save(model.state_dict(), 'model_weights.pth')

# ── 加载 ─────────────────────────────────────────────────────
model = ConvNet(num_classes=10)             # 先实例化模型
model.load_state_dict(
    torch.load('model_weights.pth',
               map_location=device)         # 指定加载到的设备
)
model.eval()   # 推理前切换到评估模式！

9.2 保存完整模型（不推荐用于生产）

# 保存整个模型（包括结构）
torch.save(model, 'full_model.pth')

# 加载（需要模型类定义在环境中）
model = torch.load('full_model.pth', map_location=device)

9.3 保存训练检查点（Checkpoint）

# 保存检查点（用于断点续训）
checkpoint = {
    'epoch':       epoch,
    'model_state': model.state_dict(),
    'optim_state': optimizer.state_dict(),
    'sched_state': scheduler.state_dict(),
    'best_acc':    best_acc,
    'config':      config,
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')

# 恢复训练
checkpoint = torch.load('checkpoint_epoch_30.pth', map_location=device)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optim_state'])
scheduler.load_state_dict(checkpoint['sched_state'])
start_epoch = checkpoint['epoch'] + 1
best_acc    = checkpoint['best_acc']

---

10. 迁移学习

10.1 迁移学习原理

迁移学习策略：

  预训练模型（ImageNet）
        ↓
  ┌─────────────────┐
  │   冻结的骨干网络  │  ← 保留通用特征（边缘、纹理、形状等）
  │   (backbone)    │
  └─────────────────┘
        ↓
  ┌─────────────────┐
  │   新的分类头     │  ← 针对新任务训练
  │   (head)        │
  └─────────────────┘

策略1 - 特征提取：冻结全部 backbone，只训练 head
策略2 - 微调：     先冻结后解冻，逐步训练 backbone（更高精度）

10.2 使用预训练模型

import torchvision.models as models
import torch.nn as nn

# ── 加载预训练模型 ──────────────────────────────────────────
# 常用模型
resnet50   = models.resnet50(weights='IMAGENET1K_V2')
efficientnet = models.efficientnet_b0(weights='IMAGENET1K_V1')
vit_b16    = models.vit_b_16(weights='IMAGENET1K_V1')
mobilenet  = models.mobilenet_v3_small(weights='IMAGENET1K_V1')

# ── 策略1：特征提取（冻结 backbone）────────────────────────
model = models.resnet50(weights='IMAGENET1K_V2')

# 冻结所有层
for param in model.parameters():
    param.requires_grad = False

# 替换分类头（解冻这一层）
num_features = model.fc.in_features  # 2048
model.fc = nn.Sequential(
    nn.Linear(num_features, 512),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(512, num_classes)      # 新类别数
)

# 只优化分类头
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

# ── 策略2：微调（解冻部分层）───────────────────────────────
model = models.resnet50(weights='IMAGENET1K_V2')

# 先冻结全部
for param in model.parameters():
    param.requires_grad = False

# 解冻最后几个 ResBlock 和分类头
for param in model.layer3.parameters():
    param.requires_grad = True
for param in model.layer4.parameters():
    param.requires_grad = True

# 替换分类头
model.fc = nn.Linear(2048, num_classes)

# 差分学习率：backbone 使用更小的 lr
optimizer = torch.optim.AdamW([
    {'params': model.layer3.parameters(), 'lr': 1e-5},
    {'params': model.layer4.parameters(), 'lr': 1e-4},
    {'params': model.fc.parameters(),     'lr': 1e-3},
])

---

11. 分布式训练

11.1 DataParallel（单机多卡，简单版）

import torch.nn as nn

# 简单包裹即可（数据并行）
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])
model = model.to('cuda:0')

⚠️ DataParallel 存在负载不均衡问题，大规模训练推荐使用 DDP。

11.2 DistributedDataParallel（DDP，推荐）

# launch_ddp.py
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def main(rank, world_size):
    # 初始化进程组
    dist.init_process_group(
        backend='nccl',    # GPU 使用 nccl，CPU 使用 gloo
        rank=rank,
        world_size=world_size
    )
    
    torch.cuda.set_device(rank)
    device = torch.device(f'cuda:{rank}')
    
    # 模型包裹 DDP
    model = ConvNet().to(device)
    model = DDP(model, device_ids=[rank])
    
    # 数据采样器（保证每个进程拿到不同数据）
    sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
    loader  = DataLoader(train_dataset, batch_size=64, sampler=sampler)
    
    for epoch in range(epochs):
        sampler.set_epoch(epoch)   # 每个 epoch 打乱不同
        train_one_epoch(model, loader, ...)
    
    dist.destroy_process_group()

# 启动方式
# torchrun --nproc_per_node=4 launch_ddp.py

---

12. 模型部署（TorchScript & ONNX）

12.1 TorchScript

TorchScript 将 Python 模型转化为 与 Python 无关的中间表示，可在 C++ 环境运行。

import torch

model = ConvNet().eval()
dummy_input = torch.randn(1, 3, 224, 224)

# ── 方式1：torch.jit.trace（追踪模式，适合无控制流）────────
traced_model = torch.jit.trace(model, dummy_input)
torch.jit.save(traced_model, 'model_traced.pt')

# ── 方式2：torch.jit.script（脚本模式，支持控制流）─────────
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'model_scripted.pt')

# ── 加载并推理 ───────────────────────────────────────────────
loaded_model = torch.jit.load('model_traced.pt')
loaded_model.eval()

with torch.no_grad():
    output = loaded_model(dummy_input)

12.2 ONNX 导出

ONNX（Open Neural Network Exchange）是跨平台的模型交换格式，可部署到 ONNX Runtime、TensorRT 等。

import torch
import torch.onnx

model = ConvNet().eval()
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    opset_version=17,
    input_names  = ['input'],
    output_names = ['output'],
    dynamic_axes = {              # 支持动态 batch size
        'input':  {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

# 使用 ONNX Runtime 推理
import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession('model.onnx',
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = sess.run(None, {'input': input_data})
print(output[0].shape)

12.3 部署方案对比

部署场景               推荐方案
─────────────────────────────────────────────────────────
Python 服务端          原生 PyTorch + Flask/FastAPI
无 Python 环境/C++     TorchScript
跨框架/跨平台          ONNX → ONNX Runtime
NVIDIA GPU 高性能      ONNX → TensorRT
移动端（iOS）          CoreML（通过 coremltools）
移动端（Android）      TorchScript + LibTorch / ONNX
嵌入式/边缘设备        量化 + TorchScript/ONNX

---

13. PyTorch 2.x 新特性：torch.compile

PyTorch 2.0 引入的 torch.compile 是最重要的新特性，通过编译优化可提速 30-100%：

import torch

model = ConvNet().to('cuda')

# 一行代码，加速推理和训练！
model = torch.compile(model)

# 不同后端
model = torch.compile(model, backend='inductor')   # 默认，最优化
model = torch.compile(model, backend='eager')      # 调试用
model = torch.compile(model, backend='aot_eager')  # 即时编译

# 不同模式
model = torch.compile(model, mode='default')         # 平衡
model = torch.compile(model, mode='reduce-overhead')  # 减少 Python 开销
model = torch.compile(model, mode='max-autotune')     # 最大性能（编译最慢）

# 完整训练循环照常使用
for inputs, targets in train_loader:
    inputs, targets = inputs.to('cuda'), targets.to('cuda')
    with autocast():
        outputs = model(inputs)   # 自动使用编译后的模型
        loss = criterion(outputs, targets)
    ...

torch.compile 加速原理：

Python 模型代码
      ↓
  TorchDynamo（图捕获）
      ↓
  AOT Autograd（提前微分）
      ↓
  TorchInductor（后端优化）
      ↓
  生成高效的 CUDA/C++ 内核

---

14. 常用工具与生态

14.1 TorchVision（计算机视觉）

import torchvision

# 预训练模型
models.resnet50, models.vgg16, models.inception_v3
models.efficientnet_b0, models.vit_b_16, models.swin_t

# 数据集
datasets.CIFAR10, datasets.ImageNet, datasets.VOCDetection

# 变换
transforms.Compose, transforms.RandomCrop, transforms.Normalize

# 工具函数
torchvision.utils.save_image(tensor, 'output.png')
torchvision.utils.make_grid(batch_images)

14.2 TensorBoard 可视化

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/experiment_1')

for epoch in range(epochs):
    train_loss, train_acc = train_one_epoch(...)
    val_loss, val_acc = evaluate(...)
    
    # 记录标量
    writer.add_scalar('Loss/train',    train_loss, epoch)
    writer.add_scalar('Loss/val',      val_loss,   epoch)
    writer.add_scalar('Accuracy/train',train_acc,  epoch)
    writer.add_scalar('Accuracy/val',  val_acc,    epoch)
    writer.add_scalar('LR',            optimizer.param_groups[0]['lr'], epoch)

# 记录模型结构
writer.add_graph(model, dummy_input)

# 记录图像
writer.add_images('train/batch', images[:8], epoch)

# 记录参数分布
for name, param in model.named_parameters():
    writer.add_histogram(f'params/{name}', param, epoch)

writer.close()

# 在终端启动
# tensorboard --logdir=runs

14.3 模型参数统计

# 统计参数量
def count_parameters(model):
    total  = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"总参数量:   {total:,}")
    print(f"可训练参数: {trainable:,}")
    print(f"冻结参数:   {total - trainable:,}")

count_parameters(model)

# 使用 torchinfo 详细统计
# pip install torchinfo
from torchinfo import summary
summary(model, input_size=(1, 3, 224, 224))

14.4 PyTorch 生态一览

核心框架
  torch         ← 核心计算
  torchvision   ← 计算机视觉
  torchaudio    ← 音频处理
  torchtext     ← 自然语言处理

训练框架
  PyTorch Lightning   ← 结构化训练框架，减少模板代码
  Hugging Face        ← NLP/CV 预训练模型
  timm                ← 图像模型库（600+ 模型）
  Detectron2          ← 目标检测（Facebook）

可视化与调试
  TensorBoard         ← 训练过程可视化
  Weights & Biases    ← 实验追踪（wandb）
  torchinfo           ← 模型结构统计

部署
  ONNX Runtime        ← 跨平台推理
  TensorRT            ← NVIDIA GPU 高性能推理
  LibTorch (C++)      ← C++ 部署
  torch.serve         ← 模型服务化

可解释性
  Captum              ← 模型归因分析

---

15. 最佳实践与调试技巧

15.1 代码规范

# ── 推荐的项目结构 ────────────────────────────────────────────
project/
├── data/
│   ├── dataset.py      # 自定义 Dataset
│   └── transforms.py   # 数据变换
├── models/
│   ├── backbone.py     # 主干网络
│   └── head.py         # 任务头
├── utils/
│   ├── metrics.py      # 评估指标
│   └── visualize.py    # 可视化工具
├── train.py            # 训练入口
├── eval.py             # 评估脚本
├── config.py           # 超参数配置
└── README.md

15.2 常见问题与解决方案

# ── 问题1：显存 OOM（Out of Memory）────────────────────────
# 解决：
# a. 减小 batch_size
# b. 使用梯度累积模拟大 batch
accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, targets) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, targets) / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# c. 使用梯度检查点（牺牲速度换显存）
from torch.utils.checkpoint import checkpoint
x = checkpoint(self.block, x)  # 重新计算前向而不存储激活值

# d. 使用混合精度（AMP）

# ── 问题2：梯度消失/爆炸 ────────────────────────────────────
# 解决：
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 梯度裁剪
# 使用 BatchNorm/LayerNorm
# 使用残差连接
# 使用 Kaiming/Xavier 初始化

# ── 问题3：训练损失不下降 ────────────────────────────────────
# 检查清单：
# □ 学习率是否合适（太大震荡，太小收敛慢）
# □ 数据是否正常（可视化几个样本）
# □ 是否调用了 optimizer.zero_grad()
# □ 标签是否正确（是类别索引还是 one-hot）
# □ 模型是否在 .train() 模式

# ── 问题4：验证精度低（过拟合）────────────────────────────────
# 解决：
# □ 增加 Dropout
# □ 增加 L2 正则化（weight_decay）
# □ 数据增强
# □ 减少模型复杂度
# □ 早停（Early Stopping）

# ── 问题5：调试技巧 ────────────────────────────────────────
# 检查中间层输出形状
print(f"Layer output shape: {x.shape}")

# 检查梯度是否正常
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: grad norm = {param.grad.norm():.4f}")

# 检查是否有 NaN
assert not torch.isnan(loss).any(), "Loss is NaN!"

15.3 性能优化技巧

# 1. 固定随机种子（可复现性）
def set_seed(seed=42):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    import numpy as np, random
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

# 2. 启用 cuDNN 加速（非确定性但更快）
torch.backends.cudnn.benchmark = True  # 输入尺寸固定时推荐开启

# 3. DataLoader 优化
loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,       # 多进程（CPU 核数的一半）
    pin_memory=True,     # 锁页内存加速 GPU 传输
    persistent_workers=True,  # 不在 epoch 结束时销毁 worker
    prefetch_factor=2    # 预取批次数
)

# 4. 使用 torch.compile（PyTorch 2.x）
model = torch.compile(model)

# 5. 推理时禁用梯度
model.eval()
with torch.no_grad():
    output = model(input)

# 6. 内存清理
del large_tensor
torch.cuda.empty_cache()

# 7. 使用 BF16（现代 GPU 推荐）
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    output = model(input)

15.4 常用代码片段速查

# ── 模型切换状态 ──────────────────────────────────────────────
model.train()   # 训练模式：启用 Dropout + BatchNorm
model.eval()    # 评估模式：禁用 Dropout，BatchNorm 使用统计量

# ── 参数冻结/解冻 ──────────────────────────────────────────────
for param in model.parameters():
    param.requires_grad = False  # 冻结
    param.requires_grad = True   # 解冻

# ── 获取当前学习率 ──────────────────────────────────────────────
current_lr = optimizer.param_groups[0]['lr']

# ── Tensor 类型转换 ────────────────────────────────────────────
t.float()    # → float32
t.half()     # → float16
t.long()     # → int64
t.bool()     # → bool
t.cpu()      # → CPU
t.cuda()     # → GPU

# ── 单个元素提取 ────────────────────────────────────────────────
scalar_tensor = torch.tensor(3.14)
value = scalar_tensor.item()   # 提取 Python 标量

# ── 批量预测 ────────────────────────────────────────────────────
model.eval()
all_preds, all_labels = [], []
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs = inputs.to(device)
        outputs = model(inputs)
        preds = outputs.argmax(dim=1).cpu().numpy()
        all_preds.extend(preds)
        all_labels.extend(labels.numpy())

---

附录：常用模块导入速查

# 核心
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import Dataset, DataLoader, random_split

# 视觉
import torchvision
import torchvision.transforms as transforms
from torchvision import datasets, models

# 自动微分
from torch.cuda.amp import autocast, GradScaler

# 分布式
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

# 可视化
from torch.utils.tensorboard import SummaryWriter

# 编译优化（PyTorch 2.x）
# model = torch.compile(model)

# 序列化
import torch.onnx
import torch.jit

---

参考资料

• PyTorch 官方文档
• PyTorch 官方教程
• PyTorch 中文文档
• 动手学深度学习（PyTorch 版）
• 深入浅出 PyTorch
• Zero to Mastery Learn PyTorch
• PyTorch Lightning 文档

---