第 8 周：PyTorch 实战

阅读指南与约定

• 所有数学公式使用 LaTeX（行内 \( ... \) 或块级 $$ ... $$）。
• 所有代码示例使用 PyTorch，并尽量给出可直接运行的最小脚本（可在 CPU 上运行；GPU/分布式部分给出可选分支）。
• 配图引用使用 Markdown 图片语法，路径形如 /sec08_xxx.png（本文提供占位引用）。
• 架构图优先使用 Mermaid（易复制到笔记/文档系统），并配合图片占位引用。

0. 环境准备（10分钟）

0.1 版本建议

• Python >= 3.10（推荐 3.11/3.12）
• PyTorch >= 2.x（含 torch.compile 可选）
• torchvision / torchaudio（按需）
• datasets（HuggingFace）
• albumentations（图像增强）

0.2 安装命令（示例）

pip install -U torch torchvision torchaudio
# 可选
pip install -U datasets albumentations opencv-python
pip install -U tensorboard wandb optuna

0.3 可复现性：随机种子与确定性

$\text{seed} \Rightarrow (\text{PRNG}_{\text{python}},\text{PRNG}_{\text{numpy}},\text{PRNG}_{\text{torch}},\text{CuDNN})$

import os, random
import numpy as np
import torch

def seed_everything(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    # 确定性（可能降低性能）
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

seed_everything(42)

8.1 PyTorch 数据管道

数据管道解决的问题：

1. 把原始数据映射为模型可消费的张量（dtype、shape、device、normalize）。
2. 在训练时高吞吐读取（多进程 worker、pin_memory、prefetch）。
3. 把样本组织成 batch（padding、对齐、变长序列）。

8.1.1 Dataset 与 DataLoader 深入

8.1.1.1 Dataset 抽象基类与契约

• __len__()：返回样本数量。
• __getitem__(idx)：返回第 idx 个样本。

形式化地，把数据集看作映射：

$\mathcal{D}: i \mapsto (x_i, y_i),\quad i\in\{0,1,\dots,N-1\}$

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __len__(self):
        raise NotImplementedError

    def __getitem__(self, idx):
        raise NotImplementedError

8.1.1.2 自定义Dataset完整实现（图片）

from pathlib import Path
from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset

class ImageFolderDataset(Dataset):
    """
    目录：
    root/
      class0/*.png
      class1/*.png
    """
    def __init__(self, root: str, transform=None):
        self.root = Path(root)
        self.transform = transform
        classes = sorted([p.name for p in self.root.iterdir() if p.is_dir()])
        self.class_to_idx = {c:i for i,c in enumerate(classes)}
        self.samples = []
        for c in classes:
            for fp in (self.root/c).glob('*'):
                if fp.suffix.lower() in {'.png','.jpg','.jpeg','.bmp'}:
                    self.samples.append((fp, self.class_to_idx[c]))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx: int):
        fp, label = self.samples[idx]
        img = Image.open(fp).convert('RGB')
        if self.transform is not None:
            x = self.transform(img)
        else:
            x = torch.from_numpy(np.array(img)).permute(2,0,1).float() / 255.0
        y = torch.tensor(label, dtype=torch.long)
        return x, y

8.1.1.3 自定义Dataset完整实现（文本）

文本工程要点：变长序列通常在 collate_fn 里 padding。

import torch
from torch.utils.data import Dataset

class CharTokenizer:
    def __init__(self):
        self.vocab = {ch:i+1 for i,ch in enumerate('abcdefghijklmnopqrstuvwxyz')}
        self.pad_id = 0
    def __call__(self, text: str):
        return [self.vocab.get(ch, 0) for ch in text.lower() if ch.isalpha()]

class ToyTextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer):
        self.texts = list(texts)
        self.labels = list(labels)
        self.tok = tokenizer
    def __len__(self):
        return len(self.texts)
    def __getitem__(self, idx):
        ids = self.tok(self.texts[idx])
        return {
            'input_ids': torch.tensor(ids, dtype=torch.long),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

padding 的数学描述：若 batch 内第 k 个序列长度 $L_k$，padding 到 $L_{\max}=\max_k L_k$：

$\tilde{x}_k=[x_k,0,0,\dots,0]\in\mathbb{Z}^{L_{\max}}$

8.1.1.4 DataLoader参数详解与性能要点

$\text{DataLoader}:(\mathcal{D},\text{sampler},\text{collate})\mapsto\{(X_b,Y_b)\}_{b=1}^B$

• batch_size：吞吐与收敛的工程折中点。
• shuffle：训练集通常 True（或显式 sampler）。
• num_workers：IO/解码瓶颈时提升明显。
• pin_memory：GPU 训练推荐 True（加速 H2D）。
• drop_last：BN/对齐需求时常用 True。

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset=ImageFolderDataset('/path/to/data'),
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=False,
)

for x, y in loader:
    pass

8.1.1.5 collate_fn：自定义批处理（变长padding）

import torch

def pad_collate(batch, pad_id: int = 0):
    input_ids = [b['input_ids'] for b in batch]
    labels = torch.stack([b['label'] for b in batch])
    lengths = torch.tensor([len(x) for x in input_ids], dtype=torch.long)
    max_len = int(lengths.max().item())

    padded = torch.full((len(batch), max_len), pad_id, dtype=torch.long)
    for i, ids in enumerate(input_ids):
        padded[i, :len(ids)] = ids

    attention_mask = (padded != pad_id).long()
    return {
        'input_ids': padded,
        'attention_mask': attention_mask,
        'lengths': lengths,
        'labels': labels,
    }

8.1.1.6 sampler策略：WeightedRandomSampler（类别不平衡）

采样概率：$P(i)=\frac{w_i}{ \sum_j w_j}$。

import torch
from torch.utils.data import DataLoader
from torch.utils.data.sampler import WeightedRandomSampler

labels = torch.tensor([0,0,0,0,1,1])
class_counts = torch.bincount(labels)
class_weights = 1.0 / class_counts.float()
weights = class_weights[labels]

sampler = WeightedRandomSampler(weights=weights, num_samples=len(weights), replacement=True)

# 这里只是演示index采样
dataset = list(range(len(labels)))
loader = DataLoader(dataset, batch_size=2, sampler=sampler)
for batch in loader:
    print(batch)

8.1.2 数据预处理与增强

增强可以视为随机变换族 $T\sim\mathcal{T}$：$\tilde{x}=T(x)$。

8.1.2.1 torchvision.transforms

from torchvision import transforms

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(0.2, 0.2, 0.2, 0.05),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize((0.4914,0.4822,0.4465), (0.247,0.243,0.261)),
])

val_tf = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914,0.4822,0.4465), (0.247,0.243,0.261)),
])

8.1.2.2 Albumentations

import albumentations as A
from albumentations.pytorch import ToTensorV2

train_aug = A.Compose([
    A.RandomResizedCrop(32, 32, scale=(0.8, 1.0)),
    A.HorizontalFlip(p=0.5),
    A.ColorJitter(p=0.5),
    A.ShiftScaleRotate(0.05, 0.1, 15, p=0.5),
    A.Normalize((0.4914,0.4822,0.4465), (0.247,0.243,0.261)),
    ToTensorV2(),
])

8.1.2.3 自定义Transform

import torch

class RandomGaussianNoise:
    def __init__(self, sigma=0.05, p=0.5):
        self.sigma = sigma
        self.p = p
    def __call__(self, x: torch.Tensor):
        if torch.rand(()) < self.p:
            x = torch.clamp(x + torch.randn_like(x) * self.sigma, 0.0, 1.0)
        return x

8.1.2.4 测试时增强（TTA）

分类常见：对 logits 的 softmax 概率平均：

$\hat{p}(y\mid x)=\frac{1}{M}\sum_{m=1}^M\text{softmax}(f(T_m(x)))$

import torch

@torch.no_grad()
def tta_predict(model, x, tta_transforms):
    model.eval()
    probs = []
    for tf in tta_transforms:
        logits = model(tf(x))
        probs.append(torch.softmax(logits, dim=-1))
    return torch.stack(probs).mean(0)

8.1.3 常用数据集加载

8.1.3.1 torchvision.datasets

from torchvision.datasets import CIFAR10, MNIST
from torchvision import transforms

root = './data'
train_ds = CIFAR10(root, train=True, download=True, transform=transforms.ToTensor())
val_ds = CIFAR10(root, train=False, download=True, transform=transforms.ToTensor())
mnist = MNIST(root, train=True, download=True, transform=transforms.ToTensor())

8.1.3.2 HuggingFace datasets

from datasets import load_dataset

ds = load_dataset('imdb')
print(ds['train'][0].keys())

8.1.3.3 自定义数据集组织方式（工程建议）

• raw/（只读原始数据）
• processed/（可再生预处理产物）
• splits/（train/val/test 切分清单）
• metadata/（label映射、统计量、hash等）

8.2 PyTorch 模型构建

8.2.1 nn.Module 深入

8.2.1.1 init 与 forward

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim: int, hidden: int, out_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )
    def forward(self, x):
        return self.net(x)

model = MLP(128, 256, 10)
print(model(torch.randn(4,128)).shape)

8.2.1.2 parameters / named_parameters / buffers

import torch
import torch.nn as nn

class WithBuffer(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(4, 4)
        self.register_buffer('scale', torch.tensor(0.5))
    def forward(self, x):
        return self.linear(x) * self.scale

m = WithBuffer()
print('buffers:', [n for n,_ in m.named_buffers()])
print('params:', [n for n,_ in m.named_parameters()])

8.2.1.3 子模块管理：add_module / children / modules

import torch
import torch.nn as nn

class Stack(nn.Module):
    def __init__(self, dims):
        super().__init__()
        for i, (a,b) in enumerate(zip(dims[:-1], dims[1:])):
            self.add_module(f'layer_{i}', nn.Linear(a,b))
    def forward(self, x):
        for layer in self.children():
            x = layer(x)
        return x

net = Stack([8,16,16,4])
print(net(torch.randn(2,8)).shape)

8.2.1.4 Hook：register_forward_hook / register_full_backward_hook

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(8,16), nn.ReLU(), nn.Linear(16,4))
acts, grads = {}, {}

handle_f = model[0].register_forward_hook(lambda m,i,o: acts.setdefault('fc1', o.detach().cpu()))

# backward hook 推荐用 full backward hook
handle_b = model[0].register_full_backward_hook(lambda m, gin, gout: grads.setdefault('fc1', gout[0].detach().cpu()))

x = torch.randn(2,8)
logits = model(x)
loss = logits.pow(2).mean()
loss.backward()

print('act:', acts['fc1'].shape)
print('grad:', grads['fc1'].shape)

handle_f.remove(); handle_b.remove()

8.2.2 常用网络层

8.2.2.1 卷积层：Conv1d/2d/3d

2D 卷积输出尺寸：