069、断点续训 Resume 源码流程：Checkpoint 的保存粒度与恢复状态机

069、断点续训 Resume 源码流程：Checkpoint 的保存粒度与恢复状态机

一、一个让我熬夜到凌晨三点的bug

去年秋天，我在调一个YOLOv8的蒸馏实验，训练到第87个epoch时，实验室突然跳闸。重启后我自信满满地加载了last.pt，结果模型loss直接飙到NaN，学习率也乱成一锅粥。更诡异的是，optimizer的momentum buffer全丢了，相当于从头开始训但学习率还卡在87epoch的位置——这种半残状态比从头训还慢。

后来我翻了一整夜源码，才发现问题出在checkpoint的保存粒度上。YOLO的断点续训远不是“存个权重就完事”那么简单，它背后藏着一个精心设计的状态机。今天我们就从源码层面，把Resume的每个齿轮都拆开看。

二、Checkpoint到底存了什么？别只盯着model

很多人以为checkpoint就是model.state_dict()，这是第一个坑。我们直接看ultralytics/engine/trainer.py里的save_checkpoint方法：

defsave_checkpoint(self):# 这里踩过坑：只存model的话，EMA权重就丢了ckpt={'epoch':self.epoch,'best_fitness':self.best_fitness,'model':deepcopy(de_parallel(self.model)).half(),# 注意是half()，省显存'ema':deepcopy(self.ema.ema).half()ifself.emaelseNone,'updates':self.ema.updatesifself.emaelseNone,'optimizer':self.optimizer.state_dict(),# 别这样写：只存param_groups'train_results':self.train_results,'wandb_id':self.wandb_run.idifself.wandb_runelseNone,'date':datetime.now().isoformat()}

注意几个关键点：

• model存的是half精度：恢复时会自动转回float32，但如果你手动load后忘了转，推理精度会掉
• optimizer存的是完整state_dict：包括momentum、learning rate scheduler的step计数。我那个bug就是因为只存了param_groups，没存state，导致Adam的momentum全清零
• ema是单独存的：很多人以为ema在model里，其实它是独立维护的滑动平均副本

保存粒度上，YOLO默认每10个epoch存一次last.pt，每100个epoch存一次best.pt。但如果你在训练中途手动中断，last.pt可能只保存了半截——比如当前epoch还没跑完，那这个checkpoint里的epoch字段就是上一个完整epoch的编号。

三、恢复状态机：从last.pt到完整训练环境

Resume的核心在resume_training方法，我们逐行拆：

defresume_training(self):# 先判断是resume还是从头训ifself.args.resume:# 这里有个隐藏逻辑：如果resume=True但没指定路径，自动找last.ptifself.args.resume=='auto':last=Path(self.args.project)/self.args.name/'weights'/'last.pt'ifnotlast.exists():# 别这样写：直接抛异常。YOLO会尝试找best.ptlast=Path(self.args.project)/self.args.name/'weights'/'best.pt'else:last=Path(self.args.resume)# 加载checkpointckpt=torch.load(last,map_location='cpu')# 关键步骤1：恢复epoch计数self.start_epoch=ckpt['epoch']+1# 注意+1，因为当前epoch已经训完self.best_fitness=ckpt['best_fitness']# 关键步骤2：恢复模型权重# 这里踩过坑：直接load_state_dict会报key不匹配，因为模型可能被DDP包装过if'model'inckpt:# 先去掉DDP的module前缀state_dict={k.replace('module.',''):vfork,vinckpt['model'].items()}self.model.load_state_dict(state_dict,strict=False)# strict=False容忍key缺失

这里有个容易忽略的细节：strict=False。为什么不用True？因为YOLO在训练过程中可能会动态添加一些层（比如自动anchor调整），这些层在checkpoint里可能没有对应key。用strict=False可以跳过这些缺失的key，但如果你改了网络结构，某些层的权重就不会被恢复——这会导致模型输出异常，但不会报错。

四、优化器恢复：最容易被忽视的坑

优化器恢复是断点续训的命门。我们继续看：

# 关键步骤3：恢复优化器状态if'optimizer'inckptandckpt['optimizer']isnotNone:# 别这样写：直接load_state_dict# 因为optimizer的param_groups顺序可能变了（比如你改了学习率策略）self.optimizer.load_state_dict(ckpt['optimizer'])# 手动恢复学习率调度器的step计数fori,groupinenumerate(self.optimizer.param_groups):if'lr'ingroup:# 这里有个隐藏逻辑：YOLO的lr是cosine衰减，需要恢复当前stepgroup['lr']=self.scheduler.get_last_lr()[i]ifhasattr(self,'scheduler')elsegroup['lr']else:# 如果没有optimizer，就重新初始化，但学习率从当前epoch开始LOGGER.warning('Optimizer not found in checkpoint, reinitializing...')self.setup_optimizer()

这个代码有个潜在问题：如果checkpoint里的optimizer和当前模型参数不匹配（比如你改了模型结构，新增了一些层），load_state_dict会报错。YOLO官方没有处理这个情况，所以如果你改了网络结构想resume，最好手动删掉optimizer字段，让代码重新初始化。

五、EMA恢复：滑动平均的精度陷阱

EMA（指数移动平均）是YOLO提升精度的关键，但恢复时容易出问题：

# 关键步骤4：恢复EMAif'ema'inckptandckpt['ema']isnotNone:# 这里踩过坑：EMA的state_dict和model的state_dict结构必须完全一致# 如果你改了模型结构，EMA恢复会失败try:self.ema.ema.load_state_dict(ckpt['ema'].float().state_dict())self.ema.updates=ckpt['updates']if'updates'inckptelse0exceptExceptionase:LOGGER.warning(f'EMA restore failed:{e}, reinitializing EMA...')self.ema=ModelEMA(self.model)else:self.ema=ModelEMA(self.model)

注意：EMA的updates字段控制着衰减系数。如果updates=0，EMA就是模型本身的副本；如果updates很大，EMA会更平滑。恢复时如果updates丢失，EMA的衰减系数会重置，导致前几个epoch的验证精度异常低——我遇到过这种情况，排查了两天才发现是EMA的decay参数不对。

六、训练状态恢复：不只是epoch

除了模型和优化器，YOLO还恢复一些“软状态”：

# 关键步骤5：恢复训练结果记录if'train_results'inckpt:self.train_results=ckpt['train_results']# 恢复最佳模型路径if'best.pt'instr(last):self.best=lastelse:self.train_results={'epoch':[],'mAP50':[],'mAP50-95':[],'loss':[]}# 关键步骤6：恢复wandb运行ID（如果有）if'wandb_id'inckptandckpt['wandb_id']:self.wandb_run=wandb.init(id=ckpt['wandb_id'],resume='must')# 关键步骤7：恢复随机数种子# 别这样写：直接设置seed。因为恢复后需要保持数据加载顺序一致ifself.args.resume:# 用当前epoch作为seed偏移，保证每个epoch的数据顺序唯一set_seed(self.args.seed+self.start_epoch)

这里有个设计哲学：YOLO的resume不是精确恢复，而是近似恢复。比如数据加载的shuffle顺序，由于随机种子变了，每个epoch的数据顺序会和中断前不同——但这不影响最终精度，只是训练曲线会有微小波动。

七、保存粒度：为什么YOLO不每epoch都存？

很多人问：为什么不每epoch都存checkpoint？这样恢复更精确。YOLO的默认策略是每10个epoch存一次，原因有三：

1. 磁盘IO开销：一个YOLOv8l的checkpoint大约500MB，每epoch存一次，100个epoch就是50GB，训练时间会多出10%
2. 恢复精度足够：丢失10个epoch的训练，对最终mAP影响通常小于0.5%，远小于超参数波动
3. 避免过拟合：频繁保存checkpoint会让开发者倾向于“选最好的checkpoint”，而不是“训到收敛”——这反而会降低泛化能力

但如果你在调参阶段，建议把save_period改成1，这样每个epoch都有checkpoint，方便回滚。代价是磁盘空间和训练时间。

八、实战经验：Resume的五个黄金法则

1. 永远不要手动修改checkpoint文件：我见过有人用文本编辑器打开.pt文件，结果文件损坏。checkpoint是二进制序列化，只能用torch.load读取
2. Resume前先验证模型结构一致性：如果你改了yaml配置文件（比如增加了检测头），resume时一定要用strict=False，并且手动检查哪些层的权重没被恢复
3. 优化器恢复失败时，不要直接忽略：如果optimizer恢复报错，建议重新初始化优化器，但把学习率设置为当前epoch对应的值——这样虽然momentum丢失，但学习率曲线是连续的
4. EMA恢复失败时，建议重新初始化：因为EMA的权重和模型权重是耦合的，如果EMA恢复失败，强行使用会导致验证精度异常
5. 多卡训练时，checkpoint只在主进程保存：DDP模式下，只有rank=0的进程会写文件。如果你在非主进程手动保存，会导致文件被覆盖

九、一个实用的Resume脚本

最后分享一个我常用的resume脚本，可以处理各种异常情况：

defsafe_resume(ckpt_path,model,optimizer,scheduler,ema=None):ckpt=torch.load(ckpt_path,map_location='cpu')# 恢复模型state_dict=ckpt.get('model',ckpt.get('ema',{}))ifstate_dict:# 去掉DDP前缀state_dict={k.replace('module.',''):vfork,vinstate_dict.items()}# 只恢复匹配的keymodel_state=model.state_dict()matched={k:vfork,vinstate_dict.items()ifkinmodel_stateandv.shape==model_state[k].shape}model.load_state_dict(matched,strict=False)print(f"Restored{len(matched)}/{len(model_state)}layers")# 恢复优化器（带异常处理）if'optimizer'inckptandoptimizer:try:optimizer.load_state_dict(ckpt['optimizer'])except:print("Optimizer restore failed, reinitializing...")# 重新初始化优化器，但保持学习率forgroupinoptimizer.param_groups:group['lr']=scheduler.get_last_lr()[0]ifschedulerelsegroup['lr']# 恢复EMAifemaand'ema'inckpt:try:ema.ema.load_state_dict(ckpt['ema'].float().state_dict())ema.updates=ckpt.get('updates',0)except:print("EMA restore failed, reinitializing...")ema=ModelEMA(model)returnckpt.get('epoch',0)+1

这个脚本的核心思想是：能恢复多少就恢复多少，但保证训练能继续。损失一些精度，总比从头训好。

十、写在最后

断点续训看似简单，实则是一个系统工程。YOLO的checkpoint设计体现了工程上的权衡：既要保证恢复的完整性，又要兼顾存储效率和代码健壮性。下次你训练到一半断电时，希望这篇文章能帮你少走弯路——至少，别再犯我那个只存param_groups的错误了。

记住：checkpoint不是存档，是状态机。每一行代码都在告诉你：训练不是线性的，而是有状态的。