mm系列权重文件瘦身

瘦身脚本:

(会在resnet50.pth文件的同级目录下生成一个resnet50_thin.pth)

import os
import torch

root_dir = os.getcwd()


def thin_pth():
    pth_path = 'resnet50_mmcls.pth'
    state_dict = torch.load(os.path.join(root_dir, pth_path))
    work_dir = os.path.dirname(pth_path)
    pth_name = str(os.path.basename(pth_path)).split('.')[0]
    out_path = os.path.join(root_dir, work_dir, f'{pth_name}_thin.pth')
    torch.save(state_dict['state_dict'], out_path)
    print("Thin pth successful")
    print('| save path:', out_path)


if __name__ == '__main__':
    thin_pth()

原理:

我们训练完的mmcls,mmdet等模型,保存好的训练权重pth是比较大的。并且里边有两个参数是我们推理的时候用不到的。

image-20220824202312137

例如mmcls中,我们训练完保存好的resnet50_mmlab.pth文件,它有三个键值对,分别是:

meta:训练时的环境配置

state_dict:我们模型的参数

optimizer:优化器的参数

我们在推理过程中其实只需要state_dict就可以了,里边保存的是我们训练好的权重参数。

瘦身结果展示:

image-20220824203036655

我们能够看到瘦身前180m,瘦身后92m,跟pytorch官方的resnet50预训练权重差不多大。


版权声明:本文为qq_39435411原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。