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是比较大的。并且里边有两个参数是我们推理的时候用不到的。

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

meta：训练时的环境配置

state_dict：我们模型的参数

optimizer：优化器的参数

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

瘦身结果展示：

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