pytorch 模型的儲存和載入
阿新 • • 發佈:2018-11-03
本文轉至:https://byjiang.com/2017/06/05/How_To_Save_And_Restore_Model/
在模型完成訓練後,我們需要將訓練好的模型儲存為一個檔案供測試使用,或者因為一些原因我們需要繼續之前的狀態訓練之前儲存的模型,那麼如何在PyTorch中儲存和恢復模型呢?
方法一(推薦):
第一種方法也是官方推薦的方法,只儲存和恢復模型中的引數。
儲存
torch.save(the_model.state_dict(), PATH)
恢復
the_model = TheModelClass(*args, **kwargs) the_model.load_state_dict(torch.load(PATH))
使用這種方法,我們需要自己匯入模型的結構資訊。
方法二:
使用這種方法,將會儲存模型的引數和結構資訊。
儲存
torch.save(the_model, PATH)
恢復
the_model = torch.load(PATH)
一個相對完整的例子
saving
torch.save({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, 'checkpoint.tar' )
loading
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict' ])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.evaluate, checkpoint['epoch']))
獲取模型中某些層的引數
對於恢復的模型,如果我們想檢視某些層的引數,可以:
# 定義一個網路
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:
params=model.state_dict()
for k,v in params.items():
print(k) #列印網路中的變數名
print(params['conv1.weight']) #列印conv1的weight
print(params['conv1.bias']) #列印conv1的bias