博客标签

标准训练流程

介绍 PyTorch 中的标准训练流程，包括训练/验证模式切换、前向传播、反向传播、优化器更新、学习率调度、AMP、梯度裁剪等内容

标准训练流程

  • 标准训练流程 = 训练 / 验证模式切换 → 取 batch → 前向 → loss → 反向 → 更新 → 清理状态 → 记录

  • 最外层结构（epoch 级）

    # epoch 是数据遍历次数
# Dataset 不知道 eposh
# Sampler（DDP）时需要知道 epoch
for epoch in range(num_epochs):
    if sampler is not None:
        sampler.set_epoch(epoch)

    train_one_epoch(...)
    val_loss = validate(...)

    if rank == 0:
        print(f"epoch {epoch}, val_loss={val_loss:.4f}")

  • train_one_epoch

    def move_to_device(batch, device):
    if isinstance(batch, torch.Tensor):
        return batch.to(device, non_blocking=True)
    if isinstance(batch, dict):
        return {k: move_to_device(v, device) for k, v in batch.items()}
    if isinstance(batch, list):
        return [move_to_device(v, device) for v in batch]
    return batch

def train_one_epoch(
    model,
    dataloader,
    optimizer,
    device,
    scheduler=None
):
    model.train() # 切换到训练模式

    # Sampler 已决定 index、DataLoader 已构造 batch、collate_fn 已对齐结构
    for batch in dataloader:
        batch = move_to_device(batch, device)

        outputs = model(**batch["inputs"])
        # loss 必须是标量，loss 必须在当前 device，loss 必须有 grad
        loss = outputs["loss"]
		# autograd 从 loss 出发，沿计算图反向传播，每个参数的 grad 被累积
        loss.backward()
        # 使用当前 grad，更新模型的参数
        optimizer.step()
        # 在更新之后，清空梯度历史，防止下一个 batch 梯度叠加
        optimizer.zero_grad()

        if scheduler is not None:
            # 学习率的时间轴，可能会是 epoch 级别的更新，即
            # for epoch:
            #     train_one_epoch()
            #     scheduler.step()
            scheduler.step()

  • validate

    def validate(model, dataloader, device):
    model.eval() # 切换至验证模式
    total_loss = 0

    with torch.no_grad(): # 无需计算梯度
        for batch in dataloader:
            batch = move_to_device(batch, device)
            outputs = model(**batch["inputs"])
            total_loss += outputs["loss"].item()

    return total_loss / len(dataloader)

  • 在一个完整的训练流程中，不应该出现 Dataset 逻辑、DataLoader 配置、Model 定义、分布式初始化

训练流程增强

  • 增强后的训练流程

    def train_one_epoch(
    model,
    dataloader,
    optimizer,
    device,
    scheduler=None,
    scaler=None,
    max_grad_norm=None
):
    model.train()

    for batch in dataloader:
        batch = move_to_device(batch, device)

        optimizer.zero_grad()

        with torch.cuda.amp.autocast(enabled=(scaler is not None)):
            outputs = model(**batch["inputs"])
            loss = outputs["loss"]

        if scaler is not None:
            scaler.scale(loss).backward()
            # AMP 的梯度裁剪
            if max_grad_norm is not None:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(),
                    max_grad_norm
                )
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            # fp32 的梯度裁剪
            if max_grad_norm is not None:
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(),
                    max_grad_norm
                )
            optimizer.step()

        if scheduler is not None:
            scheduler.step()

  • 增强后的验证流程

    def validate(model, dataloader, device):
    model.eval()
    total_loss = 0

    with torch.no_grad():
        for batch in dataloader:
            batch = move_to_device(batch, device)
            outputs = model(**batch["inputs"])
            total_loss += outputs["loss"].item()

    return total_loss / len(dataloader)

AMP

  • AMP（Automatic Mixed Precision） 要实现的是：用更低精度算得更快，但又不把数值稳定性搞崩

  • 深度学习里的计算任务可以分为两类

    • 计算密集型：matmul、conv、attention
    • 数值敏感型：loss、梯度累积、optimizer 更新

  • AMP 的核心思想是不同类型的算子，使用不同精度

  • fp16 的致命缺陷是表示范围小、精度低，在反向传播中梯度往往非常小（10−810^{-8}、10−1010^{-10}），而 fp16 表示不了 → underflow → 梯度直接变 0，从而导致梯度消失，模型无法学习

  • AMP 由两个组件构成

    • autocast → 控制 forward 的计算精度，只影响 forward，不改变模型参数的存储精度（仍是 fp32），而是在 forward 时临时切换算子精度

      # forward 才是计算密集区，backward 不在 autocast 控制范围
# 使用 BF16 训练更稳定，autocast(dtype=torch.bfloat16)
# bf16 表示范围 ≈ fp32，精度稍低但不容易 underflow，通常不需要 GradScaler
with torch.cuda.amp.autocast():
    outputs = model(x)
    loss = criterion(outputs, y)
# matmul / conv      → fp16 / bf16
# softmax / norm     → fp32
# loss               → fp32

    • GradScaler → 控制 backward 的数值稳定性，解决梯度 underflow，其先把 loss 放大，再反向传播，反向传播后再把梯度缩小

      # 初始化 scaler，scaler 内部维护当前 scale 值、是否出现 inf / nan、动态调整策略
scaler = torch.cuda.amp.GradScaler()

# backward 前先放大 loss
# 等价于 (loss * scale).backward()
scaler.scale(loss).backward()

# backward 之后缩小梯度
# 把梯度除以 scale，写回 parameter.grad
scaler.unscale_(optimizer)

# 通过 scaler 进行 optimizer.step
# 如果梯度有 inf / nan → 跳过 step
# 否则 → 正常 optimizer.step()
scaler.step(optimizer)

# 更新 scaler
# 若多次稳定 → 增大 scale
# 若出现溢出 → 减小 scale
scaler.update()

  • AMP 的完整顺序

    with torch.cuda.amp.autocast():
    outputs = model(...)
    loss = outputs["loss"]

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
# clip 需要真实梯度，必须放在 unscale 之后
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()

  • AMP 并不会减少显存里的参数大小，其参数仍是 fp32，而显存节省主要来自激活 / 中间结果

  • AMP 不保证一定更快，当 batch 很小或者 CPU/IO 瓶颈时收益有限

  • 在 loss 本身极不稳定或自定义 CUDA op 不支持 fp16 时，不能使用 AMP

  • AMP + DDP 为什么不会冲突

    • DDP 的 AllReduce 用的是 parameter.grad，而 AMP 在 AllReduce 之前已经把它 unscale 成 fp32

    • 一轮 iteration 的完整顺序

      forward (autocast)
↓
loss (fp32)
↓
scaler.scale(loss).backward()
↓
grad still scaled (临时状态)
↓
scaler.unscale_(optimizer)
↓
parameter.grad = fp32 的真实梯度
↓
DDP AllReduce (fp32 梯度)
↓
optimizer.step()

    • DDP 同步发生在 backward 结束之后，且看到的是已经 unscale 的梯度

    • 在 PyTorch 中，DDP 会在 backward 时给每个参数注册一个 autograd hook，当梯度计算完成时，hook 触发，对该参数的梯度做 AllReduce

    • 但这个 hook 看到的是 grad 张量本身，而 AMP 的流程是 backward 结束、scaler.unscale_() 改写 grad 的数值，此时 DDP hook 同步的是 unscale 后的结果

  • GradScaler 不会影响多卡一致性，因为所有 rank 用的是同一个 loss scale 策略，如果某个 rank 出现 inf /nan，scaler.step () 会在所有 rank 跳过更新

  • 推理阶段的混合精度与训练阶段的 AMP

    • 训练的 AMP 是 “数值安全工程”，而推理的混合精度是 “纯性能工程”

    • 推理阶段的流程

      with torch.autocast(device_type="cuda", dtype=torch.float16):
    y = model(x)
# 更激进地
model.half()

    • 推理阶段无需计算梯度，数值问题主要来自 softmax /normalization，而这些算子会由 autocast 保留为 fp32，因此不需要 GradScaler

    • 对比

      # 训练阶段
参数        → fp32
forward     → mixed (fp16 / fp32)
backward    → fp32
optimizer   → fp32
需要 scaler → 是
# 推理阶段
参数        → fp16 或 bf16
forward     → 低精度为主
backward    → 不存在
需要 scaler → 否

学习率调度（Scheduler）

  • Step / Cosine

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=total_steps
)
# 放在 optimizer.step() 之后，每个 step 调用一次
# scheduler 依赖当前 step 的 lr 状态，不能在 backward 前

  • Warmup 是 scheduler 的一部分，不应散落在训练循环各处

    # 前 N 个 step 线性上升学习率
# 后续交给 Cosine / Step
if global_step < warmup_steps:
    lr = base_lr * global_step / warmup_steps
    set_lr(optimizer, lr)
else:
    scheduler.step()

Checkpoint

  • 保存 checkpoint （只在 rank 0）：保存参数状态、优化器动量、scheduler 位置、AMP 缩放状态、当前 epoch

    def save_checkpoint(path, model, optimizer, scheduler, scaler, epoch):
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict() if scheduler else None,
        "scaler": scaler.state_dict() if scaler else None,
        "epoch": epoch
    }, path)

  • 恢复 checkpoint

    # load 在 model.to(device) 之后
# optimizer / scaler 必须与当初结构一致
ckpt = torch.load(path, map_location="cpu")

model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])

if scheduler and ckpt["scheduler"]:
    scheduler.load_state_dict(ckpt["scheduler"])

if scaler and ckpt["scaler"]:
    scaler.load_state_dict(ckpt["scaler"])

start_epoch = ckpt["epoch"] + 1

梯度裁剪

  • 位置放在 backward 之后，optimizer.step 之前

    # fp32
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
optimizer.step()
# AMP
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(...)
scaler.step(optimizer)

  • clip 是训练稳定性工具，不改变 loss，只限制更新幅度

早停（Early Stopping）

  • 早停不在 train_one_epoch 里，而在 epoch 外

    best_loss = float("inf")
patience = 5
bad_epochs = 0

for epoch in range(num_epochs):
    train_one_epoch(...)
    val_loss = validate(...)

    if val_loss < best_loss:
        best_loss = val_loss
        bad_epochs = 0
        save_checkpoint(...)
    else:
        bad_epochs += 1

    if bad_epochs >= patience:
        print("early stopping")
        break

完整示例

  • 一轮训练 iteration 中，真实发生的顺序是：

    • Dataset 定义「样本是什么」
    • DataLoader 定义「batch 如何产生」
    • DistributedSampler 决定「每张卡拿哪些样本」
    • model = DDP (model)
    • forward 在 autocast 里
    • loss 计算（fp32）
    • scaler.scale (loss).backward ()
    • scaler.unscale_(optimizer)
    • 可选：gradient clipping
    • optimizer.step ()

  • 第一步，实现 Dataset

    import torch
from torch.utils.data import Dataset

class ToyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data          # Tensor or numpy
        self.labels = labels

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

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.labels[idx]
        return x, y

  • 第二步，初始化分布式

    import torch
import torch.distributed as dist

def setup_ddp():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

  • 第三步，实现 DataLoader

    from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

dataset = ToyDataset(data, labels)

sampler = DistributedSampler(
    dataset,
    shuffle=True,     # 注意：shuffle 在 sampler，不在 DataLoader
    drop_last=True
)

dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True
)

  • 第四步，实现模型

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

class SimpleMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_classes):
        super().__init__()

        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, num_classes)

        self.norm = nn.LayerNorm(hidden_dim)

    def forward(self, x):
        # x: [B, D]
        x = self.fc1(x)
        x = F.relu(x)

        x = self.fc2(x)
        x = self.norm(x)
        x = F.relu(x)

        x = self.fc3(x)
        return x

  • 第五步，初始化相关组件

    import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

model = MyModel().cuda()
model = DDP(model, device_ids=[local_rank])

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

scaler = torch.cuda.amp.GradScaler()

  • 第六步，实现训练流程

    from torch.cuda.amp import autocast

def train_one_epoch(
    model,
    dataloader,
    optimizer,
    scaler,
    loss_fn,
    device,
    epoch,
    grad_accum_steps
):
    model.train()
    optimizer.zero_grad(set_to_none=True)

    total_loss = 0.0

    for step, (x, y) in enumerate(dataloader):
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)

        with autocast(device_type="cuda", dtype=torch.float16):
            logits = model(x)
            loss = loss_fn(logits, y)
            loss = loss / grad_accum_steps   # 关键：缩放 loss

        scaler.scale(loss).backward()

        is_accum_step = (step + 1) % grad_accum_steps == 0

        if is_accum_step:
            # 反向传播已经结束，开始处理真实梯度
            scaler.unscale_(optimizer)

            torch.nn.utils.clip_grad_norm_(
                model.parameters(),
                max_norm=1.0
            )

            scaler.step(optimizer)
            scaler.update()

            optimizer.zero_grad(set_to_none=True)

        total_loss += loss.item() * grad_accum_steps

    return total_loss / len(dataloader)

  • 第七步，实现验证流程

    @torch.no_grad()
def validate(
    model,
    dataloader,
    loss_fn,
    device
):
    model.eval()

    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    with autocast(device_type="cuda", dtype=torch.float16):
        for x, y in dataloader:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            logits = model(x)
            loss = loss_fn(logits, y)

            total_loss += loss.item() * x.size(0)

            preds = logits.argmax(dim=-1)
            total_correct += (preds == y).sum().item()
            total_samples += x.size(0)

    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples

    return avg_loss, accuracy

  • 第八步，完整代码流程

    def main():
    local_rank = setup_ddp()
    device = torch.device("cuda", local_rank)

    # --------------------
    # 超参数
    # --------------------
    input_dim = 128
    hidden_dim = 256
    num_classes = 10
    batch_size = 32
    grad_accum_steps = 4
    num_epochs = 10
    lr = 3e-4

    # --------------------
    # 构造假数据（示例）
    # --------------------
    num_samples = 10000
    data = torch.randn(num_samples, input_dim)
    labels = torch.randint(0, num_classes, (num_samples,))

    # --------------------
    # Dataset / Sampler / DataLoader
    # --------------------
    train_dataset = ToyDataset(data, labels)
    val_dataset = ToyDataset(data[:2000], labels[:2000])

    train_sampler = DistributedSampler(
        train_dataset,
        shuffle=True,
        drop_last=True
    )

    val_sampler = DistributedSampler(
        val_dataset,
        shuffle=False
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        sampler=train_sampler,
        num_workers=4,
        pin_memory=True,
        persistent_workers=True
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        sampler=val_sampler,
        num_workers=4,
        pin_memory=True
    )

    # --------------------
    # 模型 / DDP
    # --------------------
    model = SimpleMLP(input_dim, hidden_dim, num_classes)
    model.to(device)
    model = DDP(model, device_ids=[local_rank])

    # --------------------
    # Optimizer / Loss / AMP
    # --------------------
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    scaler = GradScaler()

    # --------------------
    # 训练循环
    # --------------------
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)

        train_loss = train_one_epoch(
            model,
            train_loader,
            optimizer,
            scaler,
            loss_fn,
            device,
            grad_accum_steps
        )

        if dist.get_rank() == 0:
            val_loss, val_acc = validate(
                model,
                val_loader,
                loss_fn,
                device
            )

            print(
                f"Epoch {epoch} | "
                f"Train Loss {train_loss:.4f} | "
                f"Val Loss {val_loss:.4f} | "
                f"Val Acc {val_acc:.4f}"
            )
这篇文档有帮助吗？
报告问题在 GitHub 上编辑

数据集

介绍 PyTorch 中数据集（Dataset）和数据加载器（DataLoader）的基本概念、设计原则、使用方法以及性能优化技巧

模型推理

介绍 PyTorch 模型推理的核心方式，包括原生模型推理和 TorchScript 推理，并对比它们的特点和适用场景，同时介绍 ONNX 和 TensorRT 两种常见的模型推理加速工具

On this page

标准训练流程
训练流程增强
AMP
学习率调度（Scheduler）
Checkpoint
梯度裁剪
早停（Early Stopping）
完整示例

页面信息