文章目录
- 方法一
- 方法二
当yolo在训练的时候,如果训练中断或者出现异常,可通过修改代码,从上一次断掉处重新训练,实现断点续训。
方法一
第一种方法:
按照官方给出的恢复训练代码,用yolo命令格式,这种情况必须是环境以安装了yolo和ultralytics两个包:
运行命令
yolo task=detect mode=train model=runs/detect/exp/weights/last.pt data=ultralytics/datasets/test.yaml epochs=100 save=True resume=True
方法二
-
在
ultralytics/yolo/engine/trainer.py
中找到check_resume
和resume_training
。 -
注释
check_resume
中resume = self.args.resume
,改成需要断点恢复的last.pt
。 -
在
resume_training
里面添加一行ckpt的值:
def check_resume(self):# resume = self.args.resume # 注释掉这一行resume = 'runs/detect/exp/weights/last.pt'; # 从最后的last.pt开始继续训练if resume:try:last = Path(check_file(resume) if isinstance(resume, (str,Path)) and Path(resume).exists() else get_latest_run())self.args = get_cfg(attempt_load_weights(last).args)self.args.model, resume = str(last), True # reinstateexcept Exception as e:raise FileNotFoundError("Resume checkpoint not found. Please pass a valid checkpoint to resume from, ""i.e. 'yolo train resume model=path/to/last.pt'") from eself.resume = resumedef resume_training(self, ckpt):ckpt = torch.load('runs/detect/exp/weights/last.pt') # 加载预训练模型if ckpt is None:returnbest_fitness = 0.0start_epoch = ckpt['epoch'] + 1if ckpt['optimizer'] is not None:self.optimizer.load_state_dict(ckpt['optimizer']) # optimizerbest_fitness = ckpt['best_fitness']if self.ema and ckpt.get('ema'):self.ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) # EMAself.ema.updates = ckpt['updates']if self.resume:assert start_epoch > 0, \f'{self.args.model} training to {self.epochs} epochs is finished, nothing to resume.\n' \f"Start a new training without --resume, i.e. 'yolo task=... mode=train model={self.args.model}'"LOGGER.info(f'Resuming training from {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs')if self.epochs < start_epoch:LOGGER.info(f"{self.model} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {self.epochs} more epochs.")self.epochs += ckpt['epoch'] # finetune additional epochsself.best_fitness = best_fitnessself.start_epoch = start_epoch
最后记住,断点续训结束后,将trainer.py还原,否则影响下次训练!!!!!!