Pytorch學(xué)習(xí) - 保存和加載模型

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

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

1. 3個(gè)函數(shù)

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

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

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

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

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

3. 保存和重載模型

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

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

（2）將整個(gè)模型保存下來(lái)，然后直接加載整個(gè)模型。（有點(diǎn)耗費(fèi)內(nèi)存…）

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

（3）如果沒(méi)有訓(xùn)練完，仍然需要繼續(xù)訓(xùn)練，除了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ù)字化實(shí)踐50位技術(shù)專家共同創(chuàng)作，文字、視頻、音頻交互閱讀

總結(jié)

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

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