pytorch 中的 state_dict 是一個簡單的python的字典對象,將每一層與它的對應參數張量建立映射關係.(如model的每一層的weights及偏置等等)
只有那些參數可以訓練的layer纔會被保存到模型的state_dict中,如卷積層,線性層等等。按理說BN是沒有參數可保存的,然而實際上在resnet中是有保存的,因爲pytorch的nn.BatchNorm2d默認affine =True,是存在映射的參數的。
優化器對象Optimizer也有一個state_dict,它包含了優化器的狀態以及被使用的超參數(如lr, momentum,weight_decay等),但是似乎很少會加載Optimizer。
參考
https://blog.csdn.net/vino_cherish/article/details/84110401
https://www.jianshu.com/p/60fc57e19615
https://www.cnblogs.com/leebxo/p/10920134.html
保存/加載模型參數
#----保存----
torch.save(model.state_dict(), 'params_name.pth') #保存的文件名後綴一般是.pt或.pth
#----加載----
model=Model() #定義模型結構
model.load_state_dict(torch.load('params_name.pth')) #加載模型參數
這種方法只保存/加載模型參數,不保存/加載模型結構,是官方推薦的方法。
就是要先自己定義模型結構,然後用上面的加載模型參數。
保存/加載模型
#----保存----
torch.save(model, 'model_name.pth')
#----加載----
model = torch.load('model_name.pth')
這種方法會同時保存和加載模型的參數和結構信息。加載時不需要自己定義模型結構,直接從預訓練模型中得到模型結構和參數。
只加載重合的部分參數
如果自己的模型跟預訓練模型只有部分層是相同的,那麼可以只加載這部分相同的參數,只要設置strict
參數爲False
來忽略那些沒有匹配到的keys
即可。
#----保存----
torch.save(modelA.state_dict(), PATH)
#----加載----
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
自己選擇要保存的參數
上面兩種方法都是很籠統的加載模型結構或者參數,如果我們想要保存一些其他的變量,可以這樣做:
#----保存----
torch.save({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
'best_prec1': best_prec1,},
'checkpoint_name.tar' )
#----加載----
checkpoint = torch.load('checkpoint_name.tar')
#按關鍵字獲取保存的參數
start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
state_dict=checkpoint['state_dict']
model=Model()#定義模型結構
model.load_state_dict(state_dict)
保存多個模型進一個文件中
既然checkpoint中可以按關鍵字保存不同的參數,那也就可以用關鍵字標識不同模型的參數來保存,然後同樣按關鍵字加載不同模型的參數。
#----保存----
torch.save({
'modelA_state_dict': modelA.state_dict(),
'modelB_state_dict': modelB.state_dict(),
'optimizerA_state_dict': optimizerA.state_dict(),
'optimizerB_state_dict': optimizerB.state_dict(),
...
}, PATH)
#----加載----
modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelAClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)
checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict']
modelB.load_state_dict(checkpoint['modelB_state_dict']
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict']
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict']
modelA.eval()
modelB.eval()
# or
modelA.train()
modelB.train()
在這裏,保存完模型後加載的時候有時會遇到CUDA out of memory的問題,我google到的解決方法是加上map_location=‘cpu’
checkpoint = torch.load(PATH,map_location='cpu')
查看模型中某些層的參數
假設模型的網絡結構如下:
# 定義一個網絡
from collections import OrderedDict
model = nn.Sequential(OrderedDict([
('conv1', nn.Conv2d(1,20,5)),
('relu1', nn.ReLU()),
('conv2', nn.Conv2d(20,64,5)),
('relu2', nn.ReLU())
]))
# 打印網絡的結構
print(model)
OUT:
Sequential (
(conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
(relu1): ReLU ()
(conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
(relu2): ReLU ()
)
查看獲取conv1的weight和bias:
model=Model() #定義模型結構
model.load_state_dict(torch.load('params_name.pth')) #加載模型參數
params=model.state_dict()
for k,v in params.items():
print(k) #打印網絡中的變量名
print(params['conv1.weight']) #打印conv1的weight
print(params['conv1.bias']) #打印conv1的bias
加載部分預訓練模型
有的時候想用預訓練模型的參數,但是自己的模型定義的層名字不同,無法直接匹配加載,此時可以遍歷state_dict,最後加載。(下面的例子是隻加載重合參數的意思,其實用strict
就行了,但主要是參考他如何遍歷state_dict
以及最後更新模型參數)
resnet152 = models.resnet152(pretrained=True) #加載模型結構和參數
pretrained_dict = resnet152.state_dict()
"""加載torchvision中的預訓練模型和參數後通過state_dict()方法提取參數
也可以直接從官方model_zoo下載:
pretrained_dict = model_zoo.load_url(model_urls['resnet152'])"""
model_dict = model.state_dict()
# 將pretrained_dict裏不屬於model_dict的鍵剔除掉
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 更新現有的model_dict
model_dict.update(pretrained_dict)
# 加載我們真正需要的state_dict
model.load_state_dict(model_dict)
或者寫詳細一點:
model_dict = model.state_dict()
state_dict = {}
for k, v in pretrained_dict.items():
if k in model_dict.keys():
# state_dict.setdefault(k, v)
state_dict[k] = v
else:
print("Missing key(s) in state_dict :{}".format(k))
model_dict.update(state_dict)
model.load_state_dict(model_dict)
加載部分預訓練模型基本都是這麼做的,這個是最詳細的:Pytorch模型遷移和遷移學習,導入部分模型參數,他還包括了修改預訓練模型參數名的方法(其實也是很簡單的):
def string_rename(old_string, new_string, start, end):
new_string = old_string[:start] + new_string + old_string[end:]
return new_string
def modify_model(pretrained_file, model, old_prefix, new_prefix):
'''
:param pretrained_file:
:param model:
:param old_prefix:
:param new_prefix:
:return:
'''
pretrained_dict = torch.load(pretrained_file)
model_dict = model.state_dict()
state_dict = modify_state_dict(pretrained_dict, model_dict, old_prefix, new_prefix)
model.load_state_dict(state_dict)
return model
def modify_state_dict(pretrained_dict, model_dict, old_prefix, new_prefix):
'''
修改model dict
:param pretrained_dict:
:param model_dict:
:param old_prefix:
:param new_prefix:
:return:
'''
state_dict = {}
for k, v in pretrained_dict.items():
if k in model_dict.keys():
# state_dict.setdefault(k, v)
state_dict[k] = v
else:
for o, n in zip(old_prefix, new_prefix):
prefix = k[:len(o)]
if prefix == o:
kk = string_rename(old_string=k, new_string=n, start=0, end=len(o))
print("rename layer modules:{}-->{}".format(k, kk))
state_dict[kk] = v
return state_dict
此外還有一個具體的例子:Pytorch自由載入部分模型參數並凍結
構建一個字典,使得字典的keys和我們自己創建的網絡相同,我們再從各種預訓練網絡把想要的參數對着新的keys填進去就可以有一個新的state_dict了,這樣我們就可以load這個新的state_dict,目前只能想到這個方法應對較爲複雜的網絡變換。
保存torch.nn.DataParallel模型
使用單GPU訓練和使用DataParallel模塊進行多GPU訓練所保存的模型有所不同,主要是DataParallel保存模型時多了一個module關鍵字,所以當使用DataParallel來加載一個不使用DataParallel訓練出來的模型時,就會報錯。
有一個技巧是:
- 如果預訓練模型是用DataParallel訓練的,那麼我們就先做
model=DataParallel(model)
然後再加載預訓練模型 - 如果預訓練模型沒有用DataParallel訓練,那麼我們就先加載預訓練模型,再做
model=DataParallel(model)
torchvision.models預訓練模型
torchvision.models這個包中包含alexnet、densenet、inception、resnet、 squeezenet、vgg等常用的網絡結構,並且提供了預訓練模型,可以通過簡單調用來讀取網絡結構和預訓練模型。
import torchvision
model = torchvision.models.resnet50(pretrained=True)#獲取網絡結構和預訓練模型
pretrained=True 會加載預訓練模型,他的加載源碼爲:
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo # model_zoo是和導入預訓練模型相關的包
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
'resnet152']
# model_urls這個字典是預訓練模型的下載地址
model_urls = {
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
若當前模型的網絡結構層與model_zoo提供的預訓練模型的網絡結構層嚴格對等:
def resnet50(pretrained=False, **kwargs):
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
網絡結構層不對等(事實上跟上面完全相同,只是多寫了幾個步驟而已):
def resnet50_cbam(pretrained=False, **kwargs):
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained:
pretrained_state_dict = model_zoo.load_url(model_urls['resnet50'])
#load_url函數根據model_urls字典下載或導入相應的預訓練模型
now_state_dict = model.state_dict() # 返回當前model模塊的字典
now_state_dict.update(pretrained_state_dict)
model.load_state_dict(now_state_dict) #最後通過調用model的load_state_dict方法用預訓練的模型參數來初始化你構建的網絡結構,這個方法就是PyTorch中通用的用一個模型的參數初始化另一個模型的層的操作。load_state_dict方法還有一個重要的參數是strict,該參數默認是True,表示預訓練模型的層和你的網絡結構層嚴格對應相等(比如層名和維度)
return model