其實之前筆者寫代碼的時候用到模型的保存和加載，需要用的時候就去度娘搜一下大致代碼，現在有時間就來整理下整個pytorch模型的保存和加載，開始學習把~
pytorch的模型和參數是分開的，可以分別保存或加載模型和參數。所以pytorch的保存和加載對應存在兩種方式：
1. 直接保存加載模型
（1）保存和加載整個模型
# 保存模型
torch.save(model, 'model.pth\pkl\pt') #一般形式torch.save(net, PATH)
# 加載模型
model = torch.load('model.pth\pkl\pt') #一般形式為model_dict=torch.load(PATH)
（2）僅保存和加載模型參數(推薦使用,需要提前手動構建模型)
速度快，占空間少
# 保存模型參數
torch.save(model.state_dict(), 'model.pth\pkl\pt') #一般形式為torch.save(net.state_dict(),PATH)
# 加載模型參數
model.load_state_dict(torch.load('model.pth\pkl\pt') #一般形式為model_dict=model.load_state_dict(torch.load(PATH))
state_dict（） 是一個Python字典，將每一層映射成它的參數張量。注意只有帶有可學習參數的層（卷積層、全連接層等），以及注冊的緩存（batchnorm的運行平均值）在state_dict 中才有記錄。state_dict同樣包含優化器對象，存儲了優化器的狀態，所使用到的超參數。
然而，在實驗中往往需要保存更多的信息，比如優化器的參數，那么可以采取下面的方法保存：
torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN,
'optimizer': optimizer.state_dict(),'alpha': loss.alpha, 'gamma': loss.gamma},
checkpoint_path + '/m-' + launchTimestamp + '-' + str("%.4f" % lossMIN) + '.pth.tar')
如下一個完整的使用model.state_dict()和optimizer.state_dict()例子：
# 定義模型
class TheModelClass(nn.Module): #定義一個神經網絡模型 TheModelClass
def __init__(self):
super(TheModelClass, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# 初始化模型
model = TheModelClass()
# 初始化優化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 打印模型的 state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict(): # param_tensor 為參數名稱
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印優化器的 state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
輸出結果：
Model's state_dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])
Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]
（3）load提供了很多重載的功能，其可以把在GPU上訓練的權重加載到CPU上跑
torch.load('tensors.pt')
# 強制所有GPU張量加載到CPU中
torch.load('tensors.pt', map_location=lambda storage, loc: storage) #或者model.load_state_dict(torch.load('model.pth', map_location='cpu'))
# 把所有的張量加載到GPU 1中
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# 把張量從GPU 1 移動到 GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
上述代碼只有在模型在一個GPU上訓練時才起作用。如果我在多個GPU上訓練模型并且保存它，然后嘗試在CPU上加載，會得到錯誤：KeyError: ‘unexpected key “module.conv1.weight” in state_dict' 如何解決？
因為此時已經使用模型保存了模型nn.DataParallel，該模型將模型存儲在該模型中module，而現在您正試圖加載模型DataParallel。您可以nn.DataParallel在網絡中暫時添加一個加載目的，也可以加載權重文件，創建一個沒有module前綴的新的有序字典，然后加載它。看第（4）點
（4）通過DataParalle使用多GPU時的保存和加載
odel=DataParalle(model)
#保存參數
torch.save(model.module.state_dict(), 'model.pth')
由此看出多個GPU時多了一個該模型中module，加載再cpu時，創建一個沒有module前綴的新的有序字典，然后加載它。
補充：一般來說，PyTorch的模型以.pt或者.pth文件格式保存。
2. 保存加載用于推理的常規Checkpoint/或繼續訓練**
checkpoint檢查點：不僅保存模型的參數，優化器參數，還有loss，epoch等（相當于一個保存模型的文件夾）
if (epoch+1) % checkpoint_interval == 0:
checkpoint = {"model_state_dict": net.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"epoch": epoch}
path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
torch.save(checkpoint, path_checkpoint)
#或者
#保存
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, PATH)
#加載
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()
# - 或者 -
model.train()
注意：
在保存用于推理或者繼續訓練的常規檢查點的時候，除了模型的state_dict之外，還必須保存其他參數。保存優化器的state_dict也非常重要，因為它包含了模型在訓練時候優化器的緩存和參數。除此之外，還可以保存停止訓練時epoch數，最新的模型損失，額外的torch.nn.Embedding層等。
要保存多個組件，則將它們放到一個字典中，然后使用torch.save()序列化這個字典。一般來說，使用.tar文件格式來保存這些檢查點。
加載各個組件，首先初始化模型和優化器，然后使用torch.load()加載保存的字典，然后可以直接查詢字典中的值來獲取保存的組件。
同樣，評估模型的時候一定不要忘了調用model.eval()。
是不是很簡單！！以上為個人經驗，希望能給大家一個參考，也希望大家多多支持腳本之家。
您可能感興趣的文章:- pytorch 搭建神經網路的實現
- 手把手教你實現PyTorch的MNIST數據集
- Python機器學習之基于Pytorch實現貓狗分類
- 總結近幾年Pytorch基于Imgagenet數據集圖像分類模型
