如何解决yolov5训练后模型权重文件大问题?

原创
01/15 14:10
阅读数 368

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 工具查看裁剪的模型是否正确:

展开阅读全文
加载中
点击引领话题📣 发布并加入讨论🔥
打赏
0 评论
0 收藏
0
分享
返回顶部
顶部