yolov5自己训练中间过程保存的模型文件大小过大原因:
如果训练未完成,yolov5m.pt 官方的模型权重文件大小为40.8MB。训练中间的模型权重文件大小为160MB,如果最终训练完成,则模型权重文件大小为40.8MB。
因为中断过训练,不会将模型中的一些不必要的内容裁剪掉了,训练代码在训练完成后将模型中的一些不必要的内容裁剪掉了,具体可以看(yolov5 7.0.7版本为例)代码:
train.py:
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 optimizers if f is best: LOGGER.info(f'\nValidating {f}...') results, _, _, _ = validate.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 at iou 0.65 single_cls=single_cls, dataloader=val_loader, save_dir=save_dir, save_json=is_coco, verbose=True, plots=plots, callbacks=callbacks, compute_loss=compute_loss) # val best model with plots if is_coco: callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
就是这个strip_optimizer(f)函数,进行的相关操作:
general.py:
def strip_optimizer(f='best.pt', s=''): # from yolov5.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 ema for k in 'optimizer', 'best_fitness', 'ema', 'updates': # keys x[k] = None x['epoch'] = -1 x['model'].half() # to FP16 for p in x['model'].parameters(): p.requires_grad = False torch.save(x, s or f) mb = os.path.getsize(s or f) / 1E6 # filesize LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
通过仿照strip_optimizer将生成模型文件的一些不必要的内容裁剪掉掉:
import os import torch def strip_optimizer(f='best.pt', s=''): # from yolov5.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 ema for k in 'optimizer', 'best_fitness', 'ema', 'updates': # keys x[k] = None x['epoch'] = -1 x['model'].half() # to FP16 for p in x['model'].parameters(): p.requires_grad = False torch.save(x, s or f) mb = os.path.getsize(s or f) / 1E6 # filesize print("Optimizer stripped from %s, saved as %s, %fMB"%(f, s if s else f, mb)) # 注意torch.save文件必须是在文件夹下, 而不能在磁盘根目录下 model_src_path = r'Z:\models\best.pt' model_dest_path = r'Z:\models\best_strip.pt' if __name__=='__main__': if os.path.exists(model_src_path): strip_optimizer(model_src_path, model_dest_path) # strip optimizers else: print("%s model file is not exist" % model_src_path)
通过Netron 工具查看裁剪的模型是否正确: