Python中torch.load()加載模型以及其map_location參數(shù)詳解
參考
torch.load()
函數(shù)格式為：torch.load(f, map_location=None, pickle_module=pickle, **pickle_load_args)，一般我們使用的時候，基本只使用前兩個參數(shù)。
模型的保存
模型保存有兩種形式，一種是保存模型的state_dict()，只是保存模型的參數(shù)。那么加載時需要先創(chuàng)建一個模型的實例model，之后通過torch.load()將保存的模型參數(shù)加載進來，得到dict，再通過model.load_state_dict(dict)將模型的參數(shù)更新。
另一種是將整個模型保存下來，之后加載的時候只需要通過torch.load()將模型加載，即可返回一個加載好的模型。
具體可參考：PyTorch模型的保存與加載。
模型加載中的map_location參數(shù)
具體來說，map_location參數(shù)是用于重定向，比如此前模型的參數(shù)是在cpu中的，我們希望將其加載到cuda:0中?；蛘呶覀冇卸鄰埧?，那么我們就可以將卡1中訓練好的模型加載到卡2中，這在數(shù)據(jù)并行的分布式深度學習中可能會用到。
首先定義一個AlexNet，并使用cuda:0將其訓練了一個貓狗分類，之后把模型存儲起來。
map_location=None
我們先把state_dict加載進來。
model_path = "./cuda_model.pth" model = torch.load(model_path) print(next(model.parameters()).device)
結(jié)果為：
cuda:0
因為保存的時候就是模型就是cuda:0的，所以加載進來也是。
map_location=torch.device()
model_path = "./cuda_model.pth"
model = torch.load(model_path, map_location=torch.device('cpu'))
print(next(model.parameters()).device)
結(jié)果為：
cpu
模型從cuda:0變成了cpu。
map_location={xx:xx}
model_path = "./cuda_model.pth"
model = torch.load(model_path, map_location={'cuda:0':'cuda:1'})
print(next(model.parameters()).device)
結(jié)果為：
cuda:1
模型從cuda:0變成了cuda:1。
model_path = "./cuda_model.pth"
model = torch.load(model_path, map_location={'cuda:2':'cpu'})
print(next(model.parameters()).device)
結(jié)果為：
cuda:0
模型還是cuda:0，并沒有變成cpu。因為這個map_location的映射是不對的，原始的模型就是cuda:0，而映射是cuda:2到cpu，是不對的。這種情況下，map_location返回None，也就是和不加map_location相同。
總結(jié)
到此這篇關于Python中torch.load()加載模型以及其map_location參數(shù)詳解的文章就介紹到這了,更多相關torch.load()加載模型map_location參數(shù)內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！
相關文章
這篇文章主要給大家介紹了關于Python中csv文件的寫入與讀取方法的相關資料,csv是"Comma-Separated Values(逗號分割的值)"的首字母縮寫,它其實和txt文件一樣,都是純文本文件,使用Python來讀寫csv文件是非常容易的,需要的朋友可以參考下2023-09-09
今天小編就為大家分享一篇基于pycharm導入模塊顯示不存在的解決方法，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2018-10-10
關于Python中的向量相加和numpy中的向量相加效率對比
今天小編就為大家分享一篇關于Python中的向量相加和numpy中的向量相加效率對比，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2019-08-08
這篇文章主要介紹了Python如何讀取相對路徑文件問題，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教2023-03-03
關于jupyter lab安裝及導入tensorflow找不到模塊的問題
這篇文章主要介紹了關于jupyter lab安裝及導入tensorflow找不到模塊的問題,本文給大家介紹的非常詳細，對大家的學習或工作具有一定的參考借鑒價值，需要的朋友可以參考下2021-03-03
這篇文章主要為大家介紹了Python CleverCSV輕松處理CSV文件全面指南,有需要的朋友可以借鑒參考下,希望能夠有所幫助,祝大家多多進步,早日升職加薪2024-01-01
這篇文章主要介紹了python進度條庫tqdm的基本操作方法，tqdm實時輸出處理進度而且占用的CPU資源非常少，支持windows、Linux、mac等系統(tǒng)，支持循環(huán)處理、多進程、遞歸處理、還可以結(jié)合linux的命令來查看處理情況等優(yōu)點，下面對其更多內(nèi)容詳細介紹,需要的朋友可以參考一下2022-03-03
這篇文章主要為大家介紹了Python光學仿真學習中對衍射算法的初步理解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進步2021-10-10
Python協(xié)程的實現(xiàn)方式小結(jié)
協(xié)程是Python中強大的并發(fā)編程工具,允許開發(fā)者編寫異步代碼以提高程序的性能和效率,在本文中,我們將深入探討Python中協(xié)程的實現(xiàn)方式,包括生成器、asyncio庫和async/await關鍵字,我們還會提供詳細的示例代碼,幫助您理解和應用協(xié)程,需要的朋友可以參考下2023-11-11
這篇文章主要為大家詳細介紹了Python中支持向量機SVM的使用方法，具有一定的參考價值，感興趣的小伙伴們可以參考一下2017-12-12