一般来说PyTorch有两种保存和读取模型参数的方法。但这篇文章我记录了一种最佳实践，可以在加载模型时避免掉一些问题。

第一种方案是保存整个模型：

1
 
torch.save(model_object, 'model.pth')

第二种方法是保存模型网络参数：

1
 
torch.save(model_object.state_dict(), 'params.pth')

加载的时候分别这样加载：

1
 
model = torch.load('model.pth')

以及：

1
 
model_object.load_state_dict(torch.load('params.pth'))

 

注意到这个方案是因为模型在加载之后，loss会飙升之后再慢慢降回来。查阅有关分析之后，判定是优化器optimizer的问题。

如果模型的保存是为了恢复训练状态，那么可以考虑同时保存优化器optimizer的参数：

1
2
3
4
5
6
7
 
state = {
    'epoch': epoch,
    'net': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    ...
}
torch.save(state, filepath)

然后这样加载：

1
2
3
4
 
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch =  checkpoint['epoch'] + 1

如果模型的保存是为了方便以后进行validation和test，可以在加载完之后制定model.eval()固定dropout和BN层。