Pytorch學習 - 保存和加載模型

• • 1. 3個函數(shù)
  • 2. 模型不同后綴名的區(qū)別
  • 3. 保存和重載模型

參考資料：
Pytorch官方文檔鏈接
某博客

1. 3個函數(shù)

torch.save() : 講一個序列化對象保存到磁盤中。使用python的pickle工具。 模型 (model)、張量 (tensor) 和各種對象的字典 (dict) 都可以用這個函數(shù)保存。

torch.load() ：將pickle對象文件反序列化到內存中，也便于將數(shù)據(jù)加載到設備中

torch.nn.Module.load_state_dict()：加載模型的參數(shù) state_dict介紹
Pytorch中，torch.nn.Module里面的可學習參數(shù)（weights和biases）都存在model.parameters()中。

2. 模型不同后綴名的區(qū)別

pytorch常見保存模型文件的后綴名有 .pt , .pth,.pkl。其實它們并不是在格式上有區(qū)別，只是后綴不同而已（僅此而已），在用torch.save()函數(shù)保存模型文件時，各人有不同的喜好，有些人喜歡用.pt后綴，有些人喜歡用.pth或.pkl.用相同的torch.save()語句保存出來的模型文件沒有什么不同。

3. 保存和重載模型

保存模型主要有兩種方式：
（1）只保存模型的參數(shù)，之后使用時再重新構建一個同樣結構的新模型，然后把保存的參數(shù)導入新模型。（推薦）

# 保存 torch.save(**model.state_dict()**,PATH) # 讀取 model = TheModelClass(*args,**kwargs) # 需要重新定義模型結構 **model.load_state_dict(torch.load(PATH))** # 根據(jù)模型結構，調用模型參數(shù) model.eval() # ??

（2）將整個模型保存下來，然后直接加載整個模型。（有點耗費內存…）

# 保存 torch.save(**model**,PATH) # 讀取 **model = torch.load(PATH)** # 不需要重構模型結構，直接load即可 model.eval()

（3）如果沒有訓練完，仍然需要繼續(xù)訓練，除了model_state_dict需要保存，還需要保存optimizer_state_dict,epoch和loss。

# 保存 torch.save({'epoch': epoch,'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss': loss,...}, PATH)# 加載 相比于前面只需要加載load_state_dict，還需要加載optimizer,epoch,loss等參數(shù) 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() # - or - model.train() 《新程序員》：云原生和全面數(shù)字化實踐50位技術專家共同創(chuàng)作，文字、視頻、音頻交互閱讀

總結

以上是生活随笔為你收集整理的Pytorch学习 - 保存模型和重新加载的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內容還不錯，歡迎將生活随笔推薦給好友。