pth 多类模型改成一类模型

import torch
import copy
def change_pth(input_pth, out_pth):
    model_dir = input_pth
    checkpoint = torch.load(model_dir) 
    model = copy.deepcopy(checkpoint)
    for name, param in model['state_dict'].items():
        if name == 'bbox_head.multi_level_conv_cls.0.weight' or name == 'bbox_head.multi_level_conv_cls.1.weight' or name == 'bbox_head.multi_level_conv_cls.2.weight':
            temp_param = param.data
            print("weight temp_param shape :", temp_param.shape)
            new_param = temp_param[:1,]
            print("weight new shape :", temp_param.shape)
            param.data = new_param
            model['state_dict'][name] = param
        if name == 'bbox_head.multi_level_conv_cls.0.bias' or name == 'bbox_head.multi_level_conv_cls.1.bias' or name == 'bbox_head.multi_level_conv_cls.2.bias':
            temp_param = param.data
            print("bias temp_param shape :", temp_param.shape)
            new_param = temp_param[:1,]
             print("bias new shape :", temp_param.shape)
            param.data = new_param
            model['state_dict'][name] = param

        print('{}, size:{}'.format(name, param.data.size()))

    torch.save(model, out_pth)