PyTorch模型的保存與加載是怎么樣的

PyTorch模型的保存與加載是怎么樣的，針對這個問題，這篇文章詳細介紹了相對應的分析和解答，希望可以幫助更多想解決這個問題的小伙伴找到更簡單易行的方法。

torch.save()和torch.load():

torch.save()和torch.load()配合使用,
分別用來保存一個對象(任何對象,
不一定要是PyTorch中的對象)到文件,和從文件中加載一個對象.
加載的時候可以指明是否需要數據在CPU和GPU中相互移動.

Module.state_dict()和Module.load_state_dict():

Module.state_dict()返回一個字典,
該字典以鍵值對的方式保存了Module的整個狀態.

Module.load_state_dict()可以從一個字典中加載參數到這個module和其后代,
如果strict是True,
那么所加載的字典和該module本身state_dict()方法返回的關鍵字必須嚴格確切的匹配上.
If strict is True, 
then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
返回值是一個命名元組:
NamedTuple with missing_keys and unexpected_keys fields,
分別保存缺失的關鍵字和未預料到的關鍵字.
如果自己的模型跟預訓練模型只有部分層是相同的，
那么可以只加載這部分相同的參數，
只要設置strict參數為False來忽略那些沒有匹配到的keys即可。

# 方式1:# model_path = 'model_name.pth'# model_params_path = 'params_name.pth'# ----保存----# torch.save(model, model_path)# ----加載----# model = torch.load(model_path)# 方式2:#----保存----# torch.save(model.state_dict(), model_params_path) #保存的文件名后綴一般是.pt或.pth
#----加載----# model=Model().cuda() #定義模型結構
# model.load_state_dict(torch.load(model_params_path))  #加載模型參數

說明:

# 保存/加載整個模型
torch.save(model, PATH)
model = torch.load(PATH)
model.eval()
這種保存/加載模型的過程使用了最直觀的語法，
所用代碼量少。這使用Python的pickle保存所有模塊。
這種方法的缺點是，保存模型的時候，
序列化的數據被綁定到了特定的類和確切的目錄。
這是因為pickle不保存模型類本身，而是保存這個類的路徑，
并且在加載的時候會使用。因此，
當在其他項目里使用或者重構的時候，加載模型的時候會出錯。




# 保存/加載 state_dict（推薦）
torch.save(model.state_dict(), PATH)
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

自己選擇要保存的參數,設置checkpoint:

#----保存----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)

保存多個模型到同一個文件:

#----保存----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')

加載預訓練模型的部分:

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] = velse:print("Missing key(s) in state_dict :{}".format(k))model_dict.update(state_dict)model.load_state_dict(model_dict)

關于PyTorch模型的保存與加載是怎么樣的問題的解答就分享到這里了，希望以上內容可以對大家有一定的幫助，如果你還有很多疑惑沒有解開，可以關注億速云行業資訊頻道了解更多相關知識。