本文轉至：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