本文主要解决模型权重迁移,主要使用pytorch读取某个权重,将其赋值给新权重格式,以下为原始代码:
顺带参数计算函数代码:
参数计算:
def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad)
权重更改代码如下:
if __name__ == '__main__': train_pth_root=r'D:\Users\User\Desktop\mask-try\epoch_100.pth' # 模型训练后得到的权重 pre_pth_root = r'D:\Users\User\Desktop\mask-try\resnet50-19c8e357.pth' # 原始预训练权重,如mmdet的resnet预训练权重 train_net=torch.load(train_pth_root) net_state_dict = train_net['state_dict'] # 训练模型权重保存字典键值 pre_net=torch.load(pre_pth_root) # 以下替换和更改成预训练权重格式,这里需要根据具体情况决定,本代码是基于mmdection修改的 keys_lst=[k.replace('backbone.','') for k in net_state_dict.keys() if 'backbone.' in k] for k,v in pre_net.items(): if k in keys_lst: k_new='backbone.'+k pre_net[k]=net_state_dict[k_new] # 保存新权重 torch.save(pre_net,'D:/Users/User/Desktop/mask-try/fasterrcnn_adaw.pth')