0.2 Python 与 PyTorch 实战：从装环境到手写训练循环

本章目标：搭建可用的 GPU 开发环境，把 PyTorch 五大件（张量/autograd/nn.Module/DataLoader/训练循环）全部打通，最后给出一份完整可运行的 CIFAR-10 训练脚本。学完标准：能独立写出一个完整的图像分类训练项目，能自己估算显存。

1. 环境（一次配好，少踩三天坑）

推荐用 uv（快）或 conda（稳）管理环境：

# 方式一：uv（推荐，快）
uv venv --python 3.11 && uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
uv pip install wandb tqdm numpy matplotlib jupyter

# 方式二：conda
conda create -n dl python=3.11 -y && conda activate dl
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

验证 GPU 可用（没有 NVIDIA 卡也没关系，本卷所有代码 CPU 能跑，只是慢；租卡见第 8 节）：

import torch
print(torch.__version__, torch.cuda.is_available())
print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU only")

2. 张量：形状操作是 80% 的日常

2.1 创建与基本属性

import torch
x = torch.randn(2, 3, 224, 224)   # 常见图像 batch：N×C×H×W
x.shape, x.dtype, x.device, x.requires_grad

2.2 图像相关：NCHW 与形状变换（必背）

图像张量约定为 NCHW（batch, channel, 高, 宽）。RGB 图在 PIL/NumPy 里通常是 HWC，进模型前要转 CHW。

img = torch.randn(224, 224, 3)          # HWC
img = img.permute(2, 0, 1)              # → CHW，注意 permute 不移动内存，只改"索引方式"
img = img.contiguous()                  # 需要真实重排内存时才调用
batch = img.unsqueeze(0)                # → NCHW，在前面加一维

reshape / view / permute 的区别（面试和 debug 双高频）：

• view：不复制数据，要求内存连续，快；
• reshape：能 view 就 view，不行就复制；
• permute/transpose：交换维度顺序，只改索引，之后常需要 contiguous。

2.3 广播（broadcasting）：省代码也埋大坑

形状不同也能算，规则是从最后一维向前对齐，缺的或为 1 的维度自动复制：

mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)  # ImageNet 通道均值
img = (img - mean) / torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

大坑：(N, 3, 224, 224) 和 (3, 224) 相加会广播成 N 份同样的操作（正确）；但 (N, 224, 224, 3) 和 (3,) 相加会把 3 广播到 W 维——shape 看似能算，语义完全错了且不报错。广播 bug 的特征：loss 能下降但结果莫名其妙。防法：关键操作前 assert x.shape == 期望形状。

2.4 dtype 与显存

dtype	字节/元素	用途
fp32	4	训练的主权重、数值基准
fp16	2	混合精度（有溢出风险，需 loss scale）
bf16	2	混合精度首选（指数位与 fp32 同宽，不易溢出，现代 GPU 默认）
int8/int4	1/0.5	推理量化

同一模型，dtype 换一半，显存减一半——这是后面所有显存估算的底层常数。

3. autograd：自动求导

w = torch.randn(3, requires_grad=True)
loss = (w * torch.tensor([1.0, 2.0, 3.0])).sum()
loss.backward()
print(w.grad)                    # = [1, 2, 3]，链式法则自动算出

三条纪律：

1. 更新参数用 with torch.no_grad(): 包住（否则梯度历史越滚越大）；
2. 每个 iteration 开头 optimizer.zero_grad()（PyTorch 默认累加梯度）；
3. 推理用 with torch.inference_mode():，比 no_grad 更省。

4. nn.Module：模型的容器

import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(3 * 32 * 32, 256), nn.ReLU(),
            nn.Linear(256, 10),
        )
    def forward(self, x):
        return self.net(x)

要点：__init__ 里赋值为 nn.Module 的属性会自动注册参数；model.state_dict() / load_state_dict() 是存取权重的标准方式（存 .pth 文件）；model.to("cuda") 把参数搬到 GPU；model.train() / model.eval() 切换 dropout/batchnorm 行为——评估忘切 eval 是经典 bug。

5. Dataset 与 DataLoader

from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms

tf = transforms.Compose([
    transforms.RandomHorizontalFlip(),          # 数据增强：只对训练集
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_set = datasets.CIFAR10("data", train=True, download=True, transform=tf)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True,
                          num_workers=4, pin_memory=True, drop_last=True)

num_workers（多进程读数据）和 pin_memory（锁页内存加速传输）是两条免费性能；drop_last=True 丢掉不满的最后一个 batch，避免训练末尾统计失真。

6. 完整训练脚本（逐段精读版）

下面是你要吃透的第一个完整项目。核心训练循环的模式，从 CIFAR 上的小 CNN 到微调 7B 的 VLM，结构完全一样：取数据 → 前向 → 算 loss → 反向 → 更新。

import torch, torch.nn as nn, time
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

torch.manual_seed(0)                                  # 可复现
device = "cuda" if torch.cuda.is_available() else "cpu"

tf_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
tf_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

train_ds = datasets.CIFAR10("data", True, download=True, transform=tf_train)
test_ds  = datasets.CIFAR10("data", False, download=True, transform=tf_test)
train_ld = DataLoader(train_ds, 256, shuffle=True,  num_workers=4, pin_memory=True, drop_last=True)
test_ld  = DataLoader(test_ds, 512, shuffle=False, num_workers=2, pin_memory=True)

class ConvBlock(nn.Sequential):
    def __init__(self, cin, cout):
        super().__init__(
            nn.Conv2d(cin, cout, 3, padding=1, bias=False),
            nn.BatchNorm2d(cout), nn.ReLU(inplace=True))

class Net(nn.Module):
    def __init__(self, ch=64):
        super().__init__()
        self.features = nn.Sequential(
            ConvBlock(3, ch),  ConvBlock(ch, ch),          # 32×32
            nn.MaxPool2d(2),                               # 16×16
            ConvBlock(ch, ch*2), ConvBlock(ch*2, ch*2),
            nn.MaxPool2d(2),                               # 8×8
            ConvBlock(ch*2, ch*4),
            nn.AdaptiveAvgPool2d(1))                       # 全局平均池化 → 1×1
        self.head = nn.Linear(ch*4, 10)
    def forward(self, x):
        x = self.features(x).flatten(1)
        return self.head(x)

model = Net().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.05)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=20)   # 20 epoch
scaler = torch.amp.GradScaler(enabled=(device == "cuda"))           # AMP 混合精度

for epoch in range(20):
    model.train(); t0 = time.time()
    for step, (x, y) in enumerate(train_ld):
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(device=="cuda")):
            logits = model(x)
            loss = nn.functional.cross_entropy(logits, y)
        opt.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.step(opt); scaler.update()
        if step % 50 == 0:
            print(f"ep{epoch} step{step} loss={loss.item():.3f} lr={sched.get_last_lr()[0]:.2e}")
    sched.step()

    model.eval(); correct = total = 0                     # 每 epoch 评估一次
    with torch.inference_mode():
        for x, y in test_ld:
            x, y = x.to(device), y.to(device)
            correct += (model(x).argmax(1) == y).sum().item(); total += y.size(0)
    print(f"[ep{epoch}] test acc={correct/total:.4f} time={time.time()-t0:.0f}s")
    torch.save(model.state_dict(), f"ckpt_ep{epoch}.pth")  # checkpoint

这个脚本里藏着本章所有知识点：种子、设备、数据增强、AMP（autocast 包前向 + GradScaler 包反向）、zero_grad、调度器、train/eval 切换、inference_mode、checkpoint。把它跑通、改坏、修好各一遍，PyTorch 就入门了。

（建议接 wandb：wandb.init(project="cifar") 后把 loss/acc wandb.log() 上去，实验管理习惯从此建立。）

7. 显存速算与 OOM 三板斧

7.1 训练显存速算（全参数、AdamW、fp16/bf16 混合精度）

$\text{显存} \approx \underbrace{2P}_{\text{混合精度权重}} + \underbrace{2P}_{\text{梯度}} + \underbrace{8P}_{\text{AdamW 两个状态(fp32)}} + \underbrace{4P}_{\text{fp32主权重}} + \text{激活值}$

其中 $P$ 是参数量。粗算约 16 字节/参数 + 激活，所以：

• 7B 模型全参数训练 ≈ 16×7GB ≈ 112GB+激活 → 单张消费级卡不可能，要 FSDP/DeepSpeed（卷7详述）；
• 7B 模型 LoRA 微调：冻结主干，只训 ~0.1% 参数 → 权重 14GB + 少量梯度和优化器状态 ≈ 24GB 卡（4090）正好能跑。这就是为什么”学生甜点区”是 LoRA + 7B/4B 模型。

7.2 OOM 三板斧（按顺序试）

1. 降 batch size，配梯度累积补等效 batch；
2. 开 gradient_checkpointing（用重算换显存，速度约慢 30%）；
3. 检查是否忘了 torch.no_grad() / inference_mode()，或把整段历史 logits 攒进了 list（日志里 loss.item() 要立刻取标量，别存 tensor）。

8. 学生算力方案

• 免费：Kaggle（每周约 30h GPU）、魔搭 ModelScope 免费算力、Colab 免费档；
• 租用：AutoDL / 仙宫云，4090 约 1.5–2 元/小时，LoRA 微调一个 7B 跑几小时就是一杯奶茶钱；
• 本地无卡：卷0、卷1 全部内容 CPU 或免费档即可完成，训练类实验攒到租卡时段集中跑。

9. 自测题

1. view 和 reshape 什么时候行为不同？（内存不连续时：view 报错，reshape 复制）
2. HWC 转 CHW 该用 view 吗？（不该，维度语义变了要用 permute；view 只能拆并相邻的内存维度）
3. optimizer.zero_grad() 忘写会怎样？（梯度跨 step 累加，loss 不降甚至爆炸）
4. eval 忘切会怎样？（dropout/BN 用训练行为，测试精度异常波动）
5. bf16 比 fp16 好在哪？（指数位同 fp32，动态范围大，不需要 loss scale，几乎不会溢出）
6. 为什么 loss.item() 而不是把 loss tensor 存进列表？（存 tensor 会保留整张计算图，显存爆）
7. batch size 减半后想保持等效更新，两个办法？（梯度累积 2 步；或调 lr）
8. 一个 3B 模型 fp16 权重占多少显存？加 AdamW 训练状态呢？（6GB；约 3×8+4=28GB 起步，还没算激活）

10. 动手作业（对应路线阶段0里程碑）

1. 跑通第 6 节脚本，记录最终精度（ConvBlock 结构约 87–91%）。
2. 把 Net 换成 torchvision.models.resnet18（num_classes=10），加更多增强（trivialaugment），目标：CIFAR-10 测试精度 ≥ 93%。这一步会逼你读懂 torchvision 源码。
3. 接入 wandb，画出 loss/lr/acc 三条曲线；故意把 lr 设成 0.1 观察发散，再设 1e-5 观察学不动——亲手建立第 4.6 节速查表的手感。
4. 手写一个自定义 Dataset：从文件夹读图片（ImageFolder 之外自己实现 __getitem__），这是后面做 VLM 数据 pipeline 的地基。

11. 延伸资源

• PyTorch 官方教程《60 Minute Blitz》→《Training a Classifier》——和本章互补
• Karpathy《Let’s build GPT from scratch》前 20 分钟——张量操作的真实工程用法
• 《Dive into Deep Learning》（动手学深度学习）中文版第 2–4 章——系统查漏