1）添加下面一句話到模型中
for p in self.parameters():
p.requires_grad = False
比如加載了resnet預訓練模型之后，在resenet的基礎上連接了新的模快，resenet模塊那部分可以先暫時凍結不更新，只更新其他部分的參數，那么可以在下面加入上面那句話
class RESNET_MF(nn.Module):
def __init__(self, model, pretrained):
super(RESNET_MF, self).__init__()
self.resnet = model(pretrained)
for p in self.parameters():
p.requires_grad = False #預訓練模型加載進來后全部設置為不更新參數，然后再后面加層
self.f = SpectralNorm(nn.Conv2d(2048, 512, 1))
self.g = SpectralNorm(nn.Conv2d(2048, 512, 1))
self.h = SpectralNorm(nn.Conv2d(2048, 2048, 1))
...
同時在優化器中添加：
filter(lambda p: p.requires_grad, model.parameters())
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)
2) 參數保存在有序的字典中，那么可以通過查找參數的名字對應的id值，進行凍結
查看每一層的代碼:
model_dict = torch.load('net.pth.tar').state_dict()
dict_name = list(model_dict)
for i, p in enumerate(dict_name):
print(i, p)
打印一下這個文件，可以看到大致是這個樣子的：
0 gamma
1 resnet.conv1.weight
2 resnet.bn1.weight
3 resnet.bn1.bias
4 resnet.bn1.running_mean
5 resnet.bn1.running_var
6 resnet.layer1.0.conv1.weight
7 resnet.layer1.0.bn1.weight
8 resnet.layer1.0.bn1.bias
9 resnet.layer1.0.bn1.running_mean
....
同樣在模型中添加這樣的代碼：
for i,p in enumerate(net.parameters()):
if i 165:
p.requires_grad = False
在優化器中添加上面的那句話可以實現參數的屏蔽
補充：pytorch 加載預訓練模型 + 斷點恢復 + 凍結訓練（避坑版本）
1、 預訓練模型網絡結構 = 你要加載模型的網絡結構
那么直接 套用
path="你的 .pt文件路徑"
model = "你的網絡"
checkpoint = torch.load(path, map_location=device)
model.load_state_dict(checkpoint)
2、 預訓練模型網絡結構 與你的網絡結構不一致
當你直接套用上面公式，會出現類似unexpected key module.xxx.weight問題
這種情況下，需要具體分析一下網絡信息，再決定如何加載。
# model_dict 是一個字典，保存網絡 各層名稱和參數，
model_dict = model.state_dict()
print(model_dict.keys()
# 這里打印出 網絡 各層名稱
checkpoint = torch.load(path,map_location=device)
for k, v in checkpoint.items():
print("keys:".k)
# 這里打印出 預訓練模型網絡 各層名稱, 是字典 【鍵】顯示的另一種方式。
然后，對比兩者網絡結構參數 的異同，
若各層網絡名稱 基本不一致，那這個預訓練模型基本就沒法用了，直接換模型吧
若兩者網絡參數有很多 類似的地方，但又不完全一致，那可以采取如下方式。
(1) 部分網絡關鍵字 ---- 完全匹配的情況
model.load_state_dict(checkpoint, strict=True)
load_state_dict 函數添加 參數 strict=True， 它直接忽略那些沒有的dict，有相同的就復制，沒有就直接放棄賦值！他要求預訓練模型的關鍵字必須確切地嚴格地和 網絡的 state_dict() 函數返回的關鍵字相匹配才能賦值。
strict 也不是很智能，適用于那些 網絡關鍵字 基本能夠匹配的情況。否則即使加載成功，網絡參數也是空的。
（2）大部分網絡關鍵字 ---- 部分匹配 （不完全相同，但類似），例如
網絡關鍵字： backbone.stage0.rbr_dense.conv.weight
預訓練模型 關鍵字：stage0.rbr_dense.conv.weight
可以看到，網絡關鍵字 比預訓練模型 多了一個前綴，其它完全一致，這種情況下，可以把 預訓練模型的 stage0.rbr_dense.conv.weight 讀入 網絡的 backbone.stage0.rbr_dense.conv.weight 中。
# 對于 字典而言，in 或 not in 運算符都是基于 key 來判斷的
model_dict = model.state_dict()
checkpoint = torch.load(path,map_location=device)
# k 是預訓練模型的一個關鍵字， ss是 網絡的有一個關鍵字
for k, v in checkpoint.items():
flag = False
for ss in model_dict.keys():
if k in ss: # 在每一個元素內部匹配
s = ss; flag = True; break
else:
continue
if flag:
checkpoint[k] = model_dict[s]
3、斷點恢復
我感覺這個和常規【模型保存加載】方法的區別主要是 epoch的恢復
# 模型保存
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
... # 有其他希望保存的內容，也可自定義
}
torch.save(state, filepath)
# 加載模型，恢復訓練
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
start_epoch = checkpoint['epoch'] + 1
4、凍結訓練
一般凍結訓練都是針對【backbone】來說的，較多應用于【遷移學習】
例如，0-49 Epoch：凍結 backbone進行訓練；50-99：不凍結訓練。
Init_Epoch = 0
Freeze_Epoch = 50
Unfreeze_Epoch =100
#------------------------------------#
# 凍結一定部分訓練
#------------------------------------#
for param in model.backbone.parameters():
param.requires_grad = False
for epoch in range(Init_Epoch,Freeze_Epoch):
# I`m Freeze-training !!
pass
#------------------------------------#
# 解凍后訓練
#------------------------------------#
for param in model.backbone.parameters():
param.requires_grad = True
for epoch in range(Freeze_Epoch,Unfreeze_Epoch):
# I`m unfreeze-training !!
pass
以上為個人經驗，希望能給大家一個參考，也希望大家多多支持腳本之家。如有錯誤或未考慮完全的地方，望不吝賜教。
您可能感興趣的文章:- Pytorch加載部分預訓練模型的參數實例
- 在pytorch中查看可訓練參數的例子
- pytorch 固定部分參數訓練的方法
- pytorch 在網絡中添加可訓練參數,修改預訓練權重文件的方法
