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)
結果為：
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)
結果為：
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)
結果為：
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)
結果為：
cuda:0
模型還是cuda:0，并沒有變成cpu。因為這個map_location的映射是不對的，原始的模型就是cuda:0，而映射是cuda:2到cpu，是不對的。這種情況下，map_location返回None，也就是和不加map_location相同。
總結
到此這篇關于Python中torch.load()加載模型以及其map_location參數(shù)詳解的文章就介紹到這了,更多相關torch.load()加載模型map_location參數(shù)內容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！
相關文章
本文介紹了Django ORM的優(yōu)缺點，然后介紹了使用原生SQL進行查詢的優(yōu)點，包括更靈活、更高效等。接著介紹了如何在Django中使用原生SQL進行查詢，包括利用Django的connection對象進行查詢以及使用Django的CursorWrapper類進行封裝。最后提醒了使用原生SQL查詢的注意事項。2023-04-04
Pandas實現(xiàn)groupby分組統(tǒng)計方法實例
在數(shù)據(jù)處理的過程,有可能需要對一堆數(shù)據(jù)分組處理,例如對不同的列進行agg聚合操作(mean，min，max等等),下面這篇文章主要給大家介紹了關于Pandas實現(xiàn)groupby分組統(tǒng)計方法的相關資料,需要的朋友可以參考下2023-06-06
python Matplotlib數(shù)據(jù)可視化（1）：簡單入門
這篇文章主要介紹了python Matplotlib的相關資料，幫助大家入門matplotlib，繪制各種圖表，感興趣的朋友可以了解下2020-09-09
Python 有很多黑魔法，為了不分你的心，今天只講 metaclass。對于 metaclass 這種特性，有兩種極端的觀點：下面小編將為大家詳細的介紹，剛興趣的小伙伴可以參考一下2021-09-09
這篇文章主要介紹了Matplotlib使用字符串代替變量繪制散點圖的方法，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧2020-02-02
Python django搭建layui提交表單,表格,圖標的實例
今天小編就為大家分享一篇Python django搭建layui提交表單,表格,圖標的實例，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2019-11-11
Python的網(wǎng)絡編程庫Gevent的安裝及使用技巧
Gevent庫的奧義在于并發(fā)式的高性能網(wǎng)絡程序設計支持,這里我們將來講解Python的網(wǎng)絡編程庫Gevent的安裝及使用技巧,來看一下Gevent支持的多進程程序編寫:2016-06-06
今天小編就為大家分享一篇Python之——生成動態(tài)路由軌跡圖的實例，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2019-11-11
這篇文章主要介紹了Python?多線程爬取案例，爬蟲屬于I/O密集型的程序，所以使用多線程可以大大提高爬取效率，文章圍繞主題展開詳細的內容介紹，具有一定的參考價值，需要的小伙伴可以參考一下2022-08-08
這篇文章主要介紹了python多線程爬取西刺代理的示例代碼，幫助大家更好的理解和學習python的爬蟲，感興趣的朋友可以了解下2021-01-01