權重pth讀取更換其值重新保存方法(附帶參數計算)

本文主要解決模型權重遷移，主要使用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')