yolov5模型训练早停模型变大

1. 背景

最近使用tph-yolov5训练yolov5l-tph-plus模型时，发现模型收敛的差不多了，就果断的停止了训练，结果发现last.pt和best.pt竟然488M，而正常训练完成的模型仅有82M.。

2. 原因分析

2.1 train代码分析

查看代码发现train.py中，训练结束后有这么一段代码：

# end training -----------------------------------------------------------------------------------------------------if RANK in [-1, 0]:LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')for f in last, best:if f.exists():strip_optimizer(f)  # strip optimizersif f is best:LOGGER.info(f'\nValidating {f}...')results, _, _ = val.run(data_dict,batch_size=batch_size // WORLD_SIZE * 2,imgsz=imgsz,model=attempt_load(f, device).half(),iou_thres=0.65 if is_coco else 0.60,  # best pycocotools results at 0.65single_cls=single_cls,dataloader=val_loader,save_dir=save_dir,save_json=is_coco,verbose=True,plots=True,callbacks=callbacks,compute_loss=compute_loss)  # val best model with plotsif is_coco:callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)callbacks.run('on_train_end', last, best, plots, epoch, results)LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")

而其中处理文件压缩的是strip_optimizer(f)

2.2 strip_optimizer函数分析

该函数位于utils/general.py中

def strip_optimizer(f='best.pt', s=''):  # from utils.general import *; strip_optimizer()# Strip optimizer from 'f' to finalize training, optionally save as 's'x = torch.load(f, map_location=torch.device('cpu'))if x.get('ema'):x['model'] = x['ema']  # replace model with emafor k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates':  # keysx[k] = Nonex['epoch'] = -1x['model'].half()  # to FP16for p in x['model'].parameters():p.requires_grad = Falsetorch.save(x, s or f)mb = os.path.getsize(s or f) / 1E6  # filesizeprint(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")

该函数用于从给定的模型文件（‘f’）中剥离优化器，并可选择性地将剥离后的模型保存为新文件（‘s’）。具体操作包括：

将模型文件加载到CPU上；
如果存在’ema’，则用’ema’替换’model’；
将’optimizer’、‘training_results’、‘wandb_id’、'ema’和’updates’这几个键的值设为None；
将’epoch’设为-1；
将模型转换为FP16；
将模型的所有参数设置为不需要梯度；
将处理后的模型保存到文件’s’中，如果’s’为空则保存到文件’f’中；
计算文件大小并输出剥离优化器后的文件名和大小。

早停没有经过该函数，因此模型精度是FP32，没有去除优化器等信息，因此模型比较大。

3. 验证

写代码调用strip_optimizer对488的模型进行处理，代码如下：

from pathlib import Path
import sys
import osFILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relativefrom utils.general import LOGGER, check_dataset, check_file, check_git_status, check_img_size, check_requirements, \check_suffix, check_yaml, colorstr, get_latest_run, increment_path, init_seeds, labels_to_class_weights, \labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizerif __name__ == '__main__':save_dir = r'E:/code/other/tph-yolov5-main/runs/train/v5l-tph-plus3/'w = save_dir + 'weights/'  # weights dirlast, best = w + 'last.pt', w + 'best.pt'for f in last, best:#if f.exists():strip_optimizer(f)

结果模型从488M变为了82M，验证成功。