首先,我们可以看到 strip_optimizer函数是在general.py文件中定义,在train.py和detect.py文件中均有调用。

定义如下:

def strip_optimizer(f='best.pt', s=''):
    """用在train.py模型训练完后 
    将optimizer、training_results、updates...从保存的模型文件f中删除
    Strip optimizer from 'f' to finalize training, optionally save as 's'
    :params f: 传入的原始保存的模型文件
    :params s: 删除optimizer等变量后的模型保存的地址 dir
    """
    # x: 为加载训练的模型
    x = torch.load(f, map_location=torch.device('cpu'))
    # 如果模型是ema replace model with ema
    if x.get('ema'):
        x['model'] = x['ema']
    # 以下模型训练涉及到的若干个指定变量置空
    for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates':  # keys
        x[k] = None
    x['epoch'] = -1  # 模型epoch恢复初始值-1
    x['model'].half()  # to FP16
    for p in x['model'].parameters():
        p.requires_grad = False
    # 保存模型 x -> s/f
    torch.save(x, s or f)
    mb = os.path.getsize(s or f) / 1E6  # filesize
    print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")

起到的作用:train.py中当训练结束时,将一些训练期间产生的模块删除,比如训练epoch为300,只有300轮结束后,才删除训练期间所需的模块,以优化模型大小。而在1-299epoch之间,也就是每个ckpt(checkpoint)模型大小为最终模型的4倍。包含着 (FP16 EMA和与模型大小相同的 FP32 优化器);detect.py中如果updata=true,则strip_optimizer函数将optimizer从ckpt中删除 更新模型(用的不多)

reference:

【YOLOV5-5.x 源码解读】general.py_满船清梦压星河HK的博客-CSDN博客

why my trained yolov5s weights is so big ?about 54.5MB · Issue #6417 · ultralytics/yolov5 · GitHubSearch before asking I have searched the YOLOv5 issues and discussions and found no similar questions. Question I use the 6.0 vision yolov5s.yaml to train with pre-trained model ,in the course of train the model weight is 54.5MB,why it i...https://github.com/ultralytics/yolov5/issues/6417

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐