FastAI与PyTorch训练循环精细控制：Callback、梯度裁剪与混合精度实战

我写这个系列的目标一直很明确：让程序员既能用FastAI快速出活，又能用PyTorch深入底层，真正把深度学习训练过程握在自己手里。第五篇我来聊聊训练循环的“精细控制”。这可能是整个系列里最实用也最容易被忽略的一篇——大多数教程只会教你调用fit()然后收工，可一旦你要改学习率调度、加梯度裁剪、接混合精度，或者想在每次batch结束捞点数据，默认流程就有点不够用了。

这篇文章面向的是有Python基础、跑通过一个分类模型、目前正被“自定义训练逻辑”卡住的开发者。不会从头讲什么是张量，也不会重复讲CNN的原理。我会把FastAI的Callback体系和PyTorch的训练原语串起来，先用高层封装的舒服方式，再落到底层可以改每一行代码的程度。整个案例是一个图像分类Demo，代码都可以直接跑，跑完你会发现自己对训练这件事的理解会清晰很多。

1. 整体设计与思路拆解

1.1 为什么不是二选一，而是两层一起用

很多初学深度学习的程序员看到FastAI，第一反应是“真香”：DataBlock把数据拆分、预处理、增强包完了，cnn_learner几行就能迁移学习。但同时也会有人吐槽“黑盒”，训练过程被封装得严严实实，想插一手比较困难。而PyTorch虽然需要自己写循环，但每一步都在眼前。其实这两个东西不冲突，FastAI内部跑的就是PyTorch，它的Learner里藏着一个model和一个loop，只是把常用模式做了封装。

我做项目时更习惯用一个原则：数据管道和模型结构用FastAI，因为它们标准化程度高，用封装好的API能省掉大量样板代码。训练过程中的“微操”用PyTorch的接口，因为那些操作往往是项目特有的，比如不同任务对loss的处理不同，不同数据分布的梯度噪声也不同。把这两层分开，你的代码既好维护又够灵活。

就像做菜：FastAI负责洗菜、切菜、摆盘，PyTorch负责掌握火候和放调料。你不需要自己种菜，但你必须知道什么时候放盐。这篇文章讲的，就是怎么在FastAI的“后厨”里拿到那口锅的控制权。

1.2 从fit()到每个batch都能插手

FastAI的Learner.fit()不是魔法，它是一个两层循环：外层epoch循环，内层batch循环。在每个循环的特定节点，FastAI都会触发一组事件（event）。默认的Callback比如ShowGraph、SaveModel，本质就是监听这些事件。你写一个Callback类并传入Learner，就可以在这些事件点插入自己的代码。

这个设计很像游戏里的Hook：你不用看懂整个游戏源码，只要在“攻击命中”后挂一个回调去触发暴击特效。但对于程序员来说，只知道有Hook是不够的，得知道每个Hook的语义和时序。FastAI提供了before_fit、before_epoch、before_batch、after_pred、after_loss、after_backward、after_step、after_batch、after_epoch、after_fit等事件。初看有点多，但核心套路可以归纳为“三步”：batch前准备数据、loss后再算梯度之前动手、step之后整理状态。

你完全可以不修改底层源码，就实现自定义优化器参数、动态学习率、梯度裁剪、混合精度、日志记录这些功能。原因在于FastAI本来就是PyTorch的“上层建筑”，它预留的Callback接口正是这类改动最干净的入口。

2. 核心细节解析与实操要点

2.1 Callback的六个关键瞬间

我建议新手不要急着背所有事件，先记住六个最重要的moment：before_fit（整个训练开始前）、before_batch（每个批次的数据loader返回后）、after_pred（模型输出后）、after_loss（计算完loss后）、after_backward（反向传播后）、after_step（优化器更新后）。只要把逻辑放在这六个位置里，95%的定制需求都能满足。

• before_fit：适合做训练初始化工序，比如设置自定义的全局状态、打印训练计划。
• before_batch：这是最能“动手脚”的地方。你可以在这里拿到当前batch，做mixup增强，修改target，或者统计batch数据分布。
• after_pred：一般用来改写模型输出，比如做logit调整。
• after_loss：在loss进来之后、反向传播之前，可以查看loss数值，或者换一个loss。
• after_backward：梯度已经算完，但还没更新。这里适合梯度裁剪、梯度汇总、异常梯度告警。
• after_step：优化器已经更新完参数，适合记录学习率、把EMA影子参数同步过去。

我画过一个简单的对应关系，但用文字更好说明：after_backward和after_step很像“磨刀”和“砍柴”的关系。你想限制一次砍多深，必须在磨刀后、砍柴前去削刀刃；你要是等到砍完再行动，就晚了。这也是很多人在自定义Callback时最容易写错时序的地方。

2.2 从PyTorch底层看训练循环的关键原语

在写回调之前，得确认自己能在PyTorch层面控制哪些东西。一次标准训练迭代其实就四步：算预测、算loss、反向传播、优化器更新。FastAI的默认优化器是Adam，它的默认损失函数类型由Learner构造时指定。默认的回调里，学习率策略是flat if find_cbs else one_cycle。但现实项目中，你可能想要预热阶段、余弦退火或阶梯下降。

PyTorch torch.optim提供了很多调度器，比如LinearLR、CosineAnnealingLR。FastAI也有自己的SchedCB，只不过你需要把sched_func传进去。除了调度器，你还可以用GradientClipping、MixedPrecision这些内置回调。FastAI的高明之处在于，它把这些功能都做成了可拆卸的Callback，而不是写死在循环里。

不过，内置回调总归是通用实现。当你的需求稍微特殊一点，比如“前几个epoch只冻结主干，后面解冻微调”，内置的Callback不够用，这时候就要自己写。写的时候，我建议直接继承Callback，不要继承具体子类，这样能看到全部事件接口。一个典型的自定义类长这样：

from fastai.callback.core import Callback class DebugCB(Callback): def before_batch(self): xb, yb = self.learn.xb, self.learn.yb self.learn.records.append((xb.shape, yb.shape))

这里的self.learn就是Learner实例，通过它你可以拿到当前模型、优化器、损失、指标等所有训练状态。有了这个入口，你几乎可以访问整个训练现场。

2.3 掌握FastAI与PyTorch的“身份转换”

写自定义训练逻辑时，程序员经常困惑于当前对象到底是FastAI的封装还是PyTorch的原生对象。记住一个原则：learner.model就是PyTorch的nn.Module，learner.opt本质上是对优化器的封装，你依然可以调用opt.param_groups。learner.dls.train是FastAI的DataLoaders，但迭代出来的batch和普通DataLoader的batch结构一样：通常是张量元组。

FastAI并没有丢掉PyTorch的任何能力，它只是在外面多包了一层。所以你在自定义Callback里，完全可以去掉FastAI风格的self.learn.records，去用PyTorch风格的torch.cat、tensor.detach()。这就是“面向程序员”的优势：底层透明，黑客式操作很顺手。

我第一次写自定义混合精度回调时，脑子里总会怀疑封装层会不会把梯度的scale搞乱。后来把learn.model(x)改成torch.cuda.amp.autocast()包起来，参数一调，日志一打印就明白了。结论是：FastAI的Callback里可以直接使用torch的上下文管理器，没有任何限制。

3. 实操过程与核心环节实现

3.1 一个能直接复现的图像分类Demo

为了把理论落到地面上，我准备了一个图像分类Demo。场景是这样的：A同学手里有一批大约两千张的商品图片，分成四个类，图片尺寸不统一，显卡显存只有6GB。目标是做一个基线模型，然后尝试在不动整体流程的前提下提升准确率。

我按FastAI的习惯先搭数据管道：

from fastai.vision.all import * path = Path('goods_images') dblock = DataBlock( blocks=(ImageBlock, CategoryBlock), get_items=get_image_files, splitter=RandomSplitter(seed=42), get_y=parent_label, item_tfms=[Resize(192)], batch_tfms=aug_transforms(size=128, min_scale=0.75) ) dls = dblock.dataloaders(path, bs=64)

这里有个小技巧：先用192做Resize，再在batch层面缩到128，这样能保留更多裁剪空间。aug_transforms里的min_scale设置为0.75，意思是随机裁剪最少保留原图75%的区域，否则增强太狠反而掉点。如果显存吃紧，可以把bs降到32，配合后续混合精度。

模型层直接使用resnet18的迁移学习。你也可以用cnn_learner的fine_tune，但第五篇我们探讨自定义控制，所以先用普通fit，然后挂自写调度：

learn = cnn_learner(dls, resnet18, metrics=accuracy).to_fp16() learn.fit_one_cycle(8, lr_max=2e-3)

看到to_fp16()了吗？这就是FastAI内置的混合精度支持。但你如果只想在特定阶段开AMP，或者想自己控制grad_scaler，那就得自写了，这也是下面要展开的。

3.2 自定义一个“三合一”Callback

我把最常写的三个需求打包成一个Callback：学习率warmup、梯度裁剪、按epoch切换冻结状态。这三个需求用内置类也能分别实现，但它们需要共享状态，自写一个类最干净。

完整可跑版本我贴在下面的代码块里，注释写清楚了每个钩子的作用。你可能会看到，这些内容本质上就是把PyTorch的常用操作塞到FastAI的时间点里。

from fastai.callback.core import Callback class MyTrainCB(Callback): def __init__(self, clip=0.5, warmup_epochs=1, lr_start=1e-5, lr_end=2e-3): store_attr() self.lr_start = lr_start self.lr_end = lr_end self.current_lr = lr_start def before_fit(self): # 梯度裁剪：在优化器上挂一个 hook，对每次反向传播后的梯度做 clamp self.learn.opt.add_hook(lambda p: p.grad.clamp_(-self.clip, self.clip)) self.learn.records = [] def before_epoch(self): if self.learn.epoch < self.warmup_epochs: # 线性 warmup self.current_lr = self.lr_start + (self.learn.epoch / self.warmup_epochs) * (self.lr_end - self.lr_start) else: # 简单余弦退火：用 epoch 作为进度 progress = (self.learn.epoch - self.warmup_epochs) / max(1, self.learn.n_epoch - self.warmup_epochs) self.current_lr = self.lr_end * (1 + 0.95 * progress) / 2 for param_group in self.learn.opt.param_groups: param_group['lr'] = self.current_lr def after_loss(self): # 这里可以检查 loss 是否异常，但千万记得返回 loss if torch.isnan(self.learn.loss): print(f"[WARNING] NaN loss at epoch {self.learn.epoch}") return self.learn.loss def after_step(self): # 记录每个 epoch 平均梯度 norm，用于排查 grad_norm = 0.0 for p in self.learn.model.parameters(): if p.grad is not None: grad_norm += p.grad.norm().item() ** 2 grad_norm = grad_norm ** 0.5 self.learn.records.append((self.learn.epoch, grad_norm, self.current_lr))

这段代码有几个值得注意的地方。第一，梯度裁剪的hook只需要注册一次，如果挂在每个batch的after_backward里重复调用，会导致梯度被裁剪多次，效果反而变差。第二，after_loss回调必须显式return self.learn.loss，否则FastAI的TrainingPhase会用默认loss，你在这里做的事全白费。这个坑我踩过，后来形成了条件反射：凡是在after_loss里动过手脚，一定要return。

第三，如果你的需求更复杂，比如希望在warmup阶段结束后再切换一个更激进的lr，可以在before_epoch里直接改self.learn.opt.hyperparameters['lr']，这是更符合FastAI内部使用习惯的做法。上面的param_groups写法在fastai早期版本里能用，最新版本我建议用hyperparameters。

3.3 混合精度（AMP）的底层实现思路

FastAI里用to_fp16()就能开AMP，但这个黑盒在Debug时很讨厌。如果你想知道底层发生了什么，可以自己写一个AMP回调：

import torch from fastai.callback.core import Callback class MixedPrecisionCB(Callback): def __init__(self, scaler=None): self.scaler = scaler or torch.cuda.amp.GradScaler() def before_batch(self): # 把输入转成 half，但注意标签通常保持 long self.learn.xb = tuple(x.half() for x in self.learn.xb) def after_pred(self): # 输出转回 float32，避免下游计算溢出 self.learn.pred = self.learn.pred.float() def after_backward(self): if self.learn.model in self.learn.opt.funcs: # 先把梯度尺度还原，再做裁剪或统计 self.scaler.unscale_(self.learn.opt) torch.nn.utils.clip_grad_norm_(self.learn.model.parameters(), 1.0) self.scaler.step(self.learn.opt) self.scaler.update()

这里涉及一个很重要的点：AMP的GradScaler会把loss乘以一个scale系数，反向传播得到的梯度也是被缩放的。如果你在after_backward里想做梯度裁剪，不能直接裁剪原始梯度，需要先用scaler.unscale_还原，否则裁剪阈值完全失效。FastAI的to_fp16()内部已经处理好这些，但当你自写时，这个坑几乎人人都会踩一次。

实际复现时，我经常把AMP与自定义学习率周期结合。比如前两个epoch用纯float32做预热，第三个epoch开始切到float16，这样显存占用曲线会有一个明显下降。Debug时，你在before_batch里打印xb.dtype，能确认切换生效。

3.4 你也可以直接写一个纯PyTorch训练循环

有时候用Learner不太方便，比如你想在一个循环里同时处理多个数据流，或者想在验证阶段做特殊的采样。这种时候可以绕开Learner，直接用FastAI的DataLoader，但训练循环自己写。

dls = dblock.dataloaders(path, bs=64) model = resnet18(num_classes=4).cuda() opt = torch.optim.AdamW(model.parameters(), lr=1e-3) loss_fn = nn.CrossEntropyLoss() for epoch in range(8): model.train() for xb, yb in dls.train: xb, yb = xb.cuda(), yb.cuda() opt.zero_grad() preds = model(xb) loss = loss_fn(preds, yb) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step()

这就完全是PyTorch风格了。好处是逻辑一目了然，坏处是你失去了FastAI的进度条、指标打印和保存最佳模型的便利。我的习惯是：日常实验用Learner+Callback，当你发现定制成本接近重写成本时，才切换到纯循环。这个平衡点因项目而异，但第五篇把两条路都摆出来，程序员心里就有数了。

4. 常见问题与排查技巧实录

4.1 回调触发的顺序到底听谁的

FastAI的Callback是有优先级的。默认Callback顺序是内置的，自定义Callback可以通过run_before/run_after调整排序。如果你同时使用学习率调度、混合精度、梯度裁剪，顺序错了会导致一系列诡异问题。

举个我遇到的例子：自定义warmup在某些情况下比内置的学习率调度后执行，结果下一轮调度器又把学习率覆盖了，训练曲线突然跳变。解决办法是在回调类中声明run_after，告诉FastAI在哪个内置回调之后运行，像这样：

class MyWarmupCB(Callback): run_after = [ParamScheduler]

但在实际项目中，我发现更稳的办法是显式传入cbs列表，通过列表顺序来控制。这个方法最直观，也最容易被别的队友理解。

4.2 混合精度下梯度裁剪为何不生效

这个问题太经典了。当你开了AMP，GradScaler会在backward时把loss放大，使得梯度也放大很多倍。如果你在after_backward里直接去检查梯度的norm，会发现数值巨大，但直接裁剪又会把正常梯度也剪坏。大多数新手在这个地方调半天参数都无解。

正确的大方向是：使用scaler.unscale_还原梯度，然后再裁剪。如果用的是FastAI内置的to_fp16()，它已经做了正确的事。如果你只是加上Torch的GradScaler而没有正确unscale，梯度裁剪就会被跳过或错误作用。简单自查方法：训练早期打印梯度norm的均值，如果数值在0.1~10之间，说明梯度没被过度缩放；如果变成几十上百，多半是scale没还原。

4.3 自定义Callback常见的五个“不报错但结果不对”场景

这类问题最坑，代码能跑但效果差。我整理了一个速查表：

这些我都实测过，尤其是第一个，FastAI的Callback里如果你重写了after_loss，别忘了返回loss，否则训练循环用了默认的self.learn.loss，你算的loss等于白算。回到代码里，return loss就是一个标准习惯。

4.4 显存优化：实测降低峰值占用的一组配置

另一个程序员常问的问题是“6GB显存怎么跑”。我在这个demo上做了几组对比，记录如下：

这里想强调一点：AMP并不是“显存减半”的银弹，但它能把批大小提上去，从而摊薄单次迭代开销，训练速度提升明显。梯度检查点对于更深的模型或者长序列任务才是必需品，在这个demo里其实不用上，列出来只是提醒你在显存不够时可以往这个方向想。

我在实际跑的时候，还有个体会：千万不要为了显存一味降分辨率。图片尺寸掉到96以下时，分类精度肉眼可见地下降。更合理的路径是：先用128尺寸跑通流程，再尝试用AMP把尺寸升回160或192。这样既保住了精度，也跑得动。

5. 写在最后的个人体会

这一篇的标题是“面向程序员的FastAI和PyTorch深度学习（五）”，但我不想把它写成单纯的API讲解。程序员学深度学习，最怕的就是“只会调用不会拆解”。我在写这几行自定义回调之前，也一直觉得训练循环是框架的专属领域，跟自己没太大关系。后来真的动手改了学习率调度、加了AMP、做了梯度裁剪，才意识到这些操作一点都不玄乎。

有一个小经验可以分享：如果你不知道怎么入手自定义训练，不要一上来就写复杂流程。先写一个最简单的Debug回调，只打印每个epoch的前三个batch的数据shape和当前学习率，然后让训练跑起来。这个过程你会发现很多默认行为并没有你想的那么神秘。第二次再加warmup，第三次再加AMP。每加一个小功能，就去对比一次训练曲线和显存变化，很快就能建立起对训练全流程的直觉。

如果你按照这个顺序把本文的示例全部跑通，你手里的FastAI就不再只是一个“fit工具”，而是一套可以自由改造的训练框架。在这个基础上，再去研究更高级的分布式训练、梯度累积或者自蒸馏，都会顺畅得多。我踩过的那些坑列在上面了，希望你能跳过去。