快捷導(dǎo)航

Python中torch.load()加載模型以及其map_location參數(shù)詳解

更新時(shí)間：2022年09月23日 09:55:41 作者：eecspan

torch.load()作用用來加載torch.save()保存的模型文件,下面這篇文章主要給大家介紹了關(guān)于Python中torch.load()加載模型以及其map_location參數(shù)的相關(guān)資料,需要的朋友可以參考下

參考

TORCH.LOAD

torch.load()

函數(shù)格式為：torch.load(f, map_location=None, pickle_module=pickle, **pickle_load_args)，一般我們使用的時(shí)候，基本只使用前兩個(gè)參數(shù)。

模型的保存

模型保存有兩種形式，一種是保存模型的state_dict()，只是保存模型的參數(shù)。那么加載時(shí)需要先創(chuàng)建一個(gè)模型的實(shí)例model，之后通過torch.load()將保存的模型參數(shù)加載進(jìn)來，得到dict，再通過model.load_state_dict(dict)將模型的參數(shù)更新。

另一種是將整個(gè)模型保存下來，之后加載的時(shí)候只需要通過torch.load()將模型加載，即可返回一個(gè)加載好的模型。

具體可參考：PyTorch模型的保存與加載。

模型加載中的map_location參數(shù)

具體來說，map_location參數(shù)是用于重定向，比如此前模型的參數(shù)是在cpu中的，我們希望將其加載到cuda:0中。或者我們有多張卡，那么我們就可以將卡1中訓(xùn)練好的模型加載到卡2中，這在數(shù)據(jù)并行的分布式深度學(xué)習(xí)中可能會(huì)用到。

首先定義一個(gè)AlexNet，并使用cuda:0將其訓(xùn)練了一個(gè)貓狗分類，之后把模型存儲(chǔ)起來。

map_location=None

我們先把state_dict加載進(jìn)來。

model_path = "./cuda_model.pth"
model = torch.load(model_path)
print(next(model.parameters()).device)

結(jié)果為：

cuda:0

因?yàn)楸４娴臅r(shí)候就是模型就是cuda:0的，所以加載進(jìn)來也是。

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。因?yàn)檫@個(gè)map_location的映射是不對(duì)的，原始的模型就是cuda:0，而映射是cuda:2到cpu，是不對(duì)的。這種情況下，map_location返回None，也就是和不加map_location相同。

總結(jié)

到此這篇關(guān)于Python中torch.load()加載模型以及其map_location參數(shù)詳解的文章就介紹到這了,更多相關(guān)torch.load()加載模型map_location參數(shù)內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: