您好,登錄后才能下訂單哦!
這篇“pytorch模型保存與加載問題怎么解決”文章的知識點大部分人都不太理解,所以小編給大家總結(jié)了以下內(nèi)容,內(nèi)容詳細,步驟清晰,具有一定的借鑒價值,希望大家閱讀完這篇文章能有所收獲,下面我們一起來看看這篇“pytorch模型保存與加載問題怎么解決”文章吧。
torch.save(model,path)
torch.load(path)
torch.save(model.state_dict(),path)
model_state_dic = torch.load(path)
model.load_state_dic(model_state_dic)
模型保存的時候會把模型結(jié)構(gòu)定義文件路徑記錄下來,加載的時候就會根據(jù)路徑解析它然后裝載參數(shù);當把模型定義文件路徑修改以后,使用torch.load(path)就會報錯。
把model文件夾修改為models后,再加載就會報錯。
import torch
from model.TextRNN import TextRNN
load_model = torch.load('experiment_model_save/textRNN.bin')
print('load_model',load_model)
這種保存完整模型結(jié)構(gòu)和參數(shù)的方式,一定不要改動模型定義文件路徑。
在多卡機器上有多張顯卡0號開始,現(xiàn)在模型在n>=1上的顯卡訓練保存后,拷貝在單卡機器上加載
import torch
from model.TextRNN import TextRNN
load_model = torch.load('experiment_model_save/textRNN_cuda_1.bin')
print('load_model',load_model)
會出現(xiàn)cuda device不匹配的問題——你保存的模代碼段 小部件型是使用的cuda1,那么采用torch.load()打開的時候,會默認的去尋找cuda1,然后把模型加載到該設備上。這個時候可以直接使用map_location來解決,把模型加載到CPU上即可。
load_model = torch.load('experiment_model_save/textRNN_cuda_1.bin',map_location=torch.device('cpu'))
當用多GPU同時訓練模型之后,不管是采用模型結(jié)構(gòu)和參數(shù)一起保存還是單獨保存模型參數(shù),然后在單卡下加載都會出現(xiàn)問題
a、模型結(jié)構(gòu)和參數(shù)一起保然后在加載
torch.distributed.init_process_group(backend='nccl')
模型訓練的時候采用上述多進程的方式,所以你在加載的時候也要聲明,不然就會報錯。
b、單獨保存模型參數(shù)
model = Transformer(num_encoder_layers=6,num_decoder_layers=6)
state_dict = torch.load('train_model/clip/experiment.pt')
model.load_state_dict(state_dict)
同樣會出現(xiàn)問題,不過這里出現(xiàn)的問題是參數(shù)字典的key和模型定義的key不一樣
原因是多GPU訓練下,使用分布式訓練的時候會給模型進行一個包裝,代碼如下:
model = torch.load('train_model/clip/Vtransformers_bert_6_layers_encoder_clip.bin')
print(model)
model.cuda(args.local_rank)
。。。。。。
model = nn.parallel.DistributedDataParallel(model,device_ids=[args.local_rank],find_unused_parameters=True)
print('model',model)
包裝前的模型結(jié)構(gòu):
包裝后的模型
在外層多了DistributedDataParallel以及module,所以才會導致在單卡環(huán)境下加載模型權重的時候出現(xiàn)權重的keys不一致。
if gpu_count > 1:
torch.save(model.module.state_dict(),save_path)
else:
torch.save(model.state_dict(),save_path)
model = Transformer(num_encoder_layers=6,num_decoder_layers=6)
state_dict = torch.load(save_path)
model.load_state_dict(state_dict)
這樣就是比較好的范式,加載不會出錯。
以上就是關于“pytorch模型保存與加載問題怎么解決”這篇文章的內(nèi)容,相信大家都有了一定的了解,希望小編分享的內(nèi)容對大家有幫助,若想了解更多相關的知識內(nèi)容,請關注億速云行業(yè)資訊頻道。
免責聲明:本站發(fā)布的內(nèi)容(圖片、視頻和文字)以原創(chuàng)、轉(zhuǎn)載和分享為主,文章觀點不代表本網(wǎng)站立場,如果涉及侵權請聯(lián)系站長郵箱:is@yisu.com進行舉報,并提供相關證據(jù),一經(jīng)查實,將立刻刪除涉嫌侵權內(nèi)容。