pytorch只加載預訓練模型中的部分參數及凍結部分參數

說明
比如我需要訓練車牌檢測模型, 採用retinanet, 結構爲bacnbone-fpn-retinanethead. 準備在coco數據集上預訓練. 但是coco數據集有81類, 車牌只有幾類. 預訓練完以後, retinanethead部分, 由於類數目尺寸不匹配, 所以希望只加載bacnbone以及fpn部分的參數.

保存的checkpoints本質上爲一個字典, 所以只需要把head部分的key, 和value去掉即可. 觀察看到retinanethead部分都含義roi_head, 所以只需要以下操作:

 model_dict=torch.load(PATH)
 new_state_dict = {}
 for k, v in state_dict.items():
    if 'roi_head' not in k:
        new_state_dict[k] = v
model.load(new_state_dict)

或者把模型保存,以後直接加載使用
torch.save(new_state_dict, ‘0.25res18-fpn-coco-pretrain.pth’)
所以只需要根據key和value選取需要的部分即可.其他同理

2.凍結部分參數
1)a)直接在模型中加入
for p in self.parameters():
p.requires_grad = False
b)
load 模型的時候, 對應的參數設爲p.requires_grad = False
2)優化器filter
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001,
betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5)

參考:https://discuss.pytorch.org/t/how-the-pytorch-freeze-network-in-some-layers-only-the-rest-of-the-training/7088
https://blog.csdn.net/qq_21997625/article/details/90369838
https://zhuanlan.zhihu.com/p/65105409