yolov5中strip_optimizer函数的作用
首先,我们可以看到 strip_optimizer函数是在general.py文件中定义,在train.py和detect.py文件中均有调用。定义如下:def strip_optimizer(f='best.pt', s=''):"""用在train.py模型训练完后将optimizer、training_results、updates...从保存的模型文件f中删除Strip optimizer
·
首先,我们可以看到 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:
更多推荐



所有评论(0)