yolov5训练模型中断后的优化压缩

在训练yolov5模型时，会遇到后续的训练已经没有意义，这时候中断训练的话，最优模型往往没有去除各种参数信息，文件比较大，参考了yolov5源代码中的utils/general.py中的strip_optimizer()函数，进行了简单的改写，可以直接用命令行完成优化。代码如下：

import torch
import os
import sys

def strip_optimizer(f='best.pt', s=''):  
    # 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(f"保存优化后的模型 {s or f}, 大小为: {mb} MB")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("使用方法: python3 optimizer_model.py <input_model.pt> [<output_model.pt>]")
        sys.exit(1)

    input_model = sys.argv[1]
    output_model = sys.argv[2] if len(sys.argv) > 2 else ''
    print("开始优化... ...")

    strip_optimizer(f=input_model, s=output_model)

如果要使用它，假设我们将它保存为optimizer_model.py，则使用如下命令：

python3 optimizer_model.py in.pt out.pt