PyTorch模型的保存與加載是怎么樣的，針對(duì)這個(gè)問題，這篇文章詳細(xì)介紹了相對(duì)應(yīng)的分析和解答，希望可以幫助更多想解決這個(gè)問題的小伙伴找到更簡單易行的方法。

成都創(chuàng)新互聯(lián)公司是專業(yè)的鄂溫克網(wǎng)站建設(shè)公司，鄂溫克接單;提供網(wǎng)站設(shè)計(jì)、成都網(wǎng)站建設(shè),網(wǎng)頁設(shè)計(jì),網(wǎng)站設(shè)計(jì),建網(wǎng)站,PHP網(wǎng)站建設(shè)等專業(yè)做網(wǎng)站服務(wù);采用PHP框架,可快速的進(jìn)行鄂溫克網(wǎng)站開發(fā)網(wǎng)頁制作和功能擴(kuò)展;專業(yè)做搜索引擎喜愛的網(wǎng)站,專業(yè)的做網(wǎng)站團(tuán)隊(duì),希望更多企業(yè)前來合作!

torch.save()和torch.load():

torch.save()和torch.load()配合使用,
分別用來保存一個(gè)對(duì)象(任何對(duì)象,
不一定要是PyTorch中的對(duì)象)到文件,和從文件中加載一個(gè)對(duì)象.
加載的時(shí)候可以指明是否需要數(shù)據(jù)在CPU和GPU中相互移動(dòng).

Module.state_dict()和Module.load_state_dict():

Module.state_dict()返回一個(gè)字典,
該字典以鍵值對(duì)的方式保存了Module的整個(gè)狀態(tài).

Module.load_state_dict()可以從一個(gè)字典中加載參數(shù)到這個(gè)module和其后代,
如果strict是True,
那么所加載的字典和該module本身state_dict()方法返回的關(guān)鍵字必須嚴(yán)格確切的匹配上.
If strict is True, 
then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.
返回值是一個(gè)命名元組:
NamedTuple with missing_keys and unexpected_keys fields,
分別保存缺失的關(guān)鍵字和未預(yù)料到的關(guān)鍵字.
如果自己的模型跟預(yù)訓(xùn)練模型只有部分層是相同的，
那么可以只加載這部分相同的參數(shù)，
只要設(shè)置strict參數(shù)為False來忽略那些沒有匹配到的keys即可。

# 方式1:# model_path = 'model_name.pth'# model_params_path = 'params_name.pth'# ----保存----# torch.save(model, model_path)# ----加載----# model = torch.load(model_path)# 方式2:#----保存----# torch.save(model.state_dict(), model_params_path) #保存的文件名后綴一般是.pt或.pth
#----加載----# model=Model().cuda() #定義模型結(jié)構(gòu)
# model.load_state_dict(torch.load(model_params_path))  #加載模型參數(shù)

說明:

# 保存/加載整個(gè)模型
torch.save(model, PATH)
model = torch.load(PATH)
model.eval()
這種保存/加載模型的過程使用了最直觀的語法，
所用代碼量少。這使用Python的pickle保存所有模塊。
這種方法的缺點(diǎn)是，保存模型的時(shí)候，
序列化的數(shù)據(jù)被綁定到了特定的類和確切的目錄。
這是因?yàn)閜ickle不保存模型類本身，而是保存這個(gè)類的路徑，
并且在加載的時(shí)候會(huì)使用。因此，
當(dāng)在其他項(xiàng)目里使用或者重構(gòu)的時(shí)候，加載模型的時(shí)候會(huì)出錯(cuò)。




# 保存/加載 state_dict（推薦）
torch.save(model.state_dict(), PATH)
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

自己選擇要保存的參數(shù),設(shè)置checkpoint:

#----保存----torch.save({
   
   
   'epoch': epoch + 1,'arch': args.arch,'state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),
  	'loss': loss,'best_prec1': best_prec1,}, 'checkpoint_name.tar' )#----加載----checkpoint = torch.load('checkpoint_name.tar')#按關(guān)鍵字獲取保存的參數(shù)
start_epoch = checkpoint['epoch']best_prec1 = checkpoint['best_prec1']state_dict=checkpoint['state_dict']model=Model()#定義模型結(jié)構(gòu)
model.load_state_dict(state_dict)

保存多個(gè)模型到同一個(gè)文件:

#----保存----torch.save({
   
   
   
  'modelA_state_dict': modelA.state_dict(),
  'modelB_state_dict': modelB.state_dict(),
  'optimizerA_state_dict': optimizerA.state_dict(),
  'optimizerB_state_dict': optimizerB.state_dict(),
  ...
  }, PATH)#----加載----modelA = TheModelAClass(*args, **kwargs)modelB = TheModelAClass(*args, **kwargs)optimizerA = TheOptimizerAClass(*args, **kwargs)optimizerB = TheOptimizerBClass(*args, **kwargs)checkpoint = torch.load(PATH)modelA.load_state_dict(checkpoint['modelA_state_dict']modelB.load_state_dict(checkpoint['modelB_state_dict']optimizerA.load_state_dict(checkpoint['optimizerA_state_dict']optimizerB.load_state_dict(checkpoint['optimizerB_state_dict']modelA.eval()modelB.eval()# or
modelA.train()modelB.train()# 在這里，保存完模型后加載的時(shí)候有時(shí)會(huì)
# 遇到CUDA out of memory的問題，
# 我google到的解決方法是加上map_location=‘cpu’

checkpoint = torch.load(PATH,map_location='cpu')

加載預(yù)訓(xùn)練模型的部分:

resnet152 = models.resnet152(pretrained=True) #加載模型結(jié)構(gòu)和參數(shù)
pretrained_dict = resnet152.state_dict()"""加載torchvision中的預(yù)訓(xùn)練模型和參數(shù)后通過state_dict()方法提取參數(shù)
   也可以直接從官方model_zoo下載：
   pretrained_dict = model_zoo.load_url(model_urls['resnet152'])"""
model_dict = model.state_dict()# 將pretrained_dict里不屬于model_dict的鍵剔除掉
pretrained_dict = {
   
   
   k: v for k, v in pretrained_dict.items() if k in model_dict}# 更新現(xiàn)有的model_dict
model_dict.update(pretrained_dict)# 加載我們真正需要的state_dict
model.load_state_dict(model_dict)

或者寫詳細(xì)一點(diǎn):

model_dict = model.state_dict()state_dict = {
   
   
   }for k, v in pretrained_dict.items():if k in model_dict.keys():# state_dict.setdefault(k, v)state_dict[k] = velse:print("Missing key(s) in state_dict :{}".format(k))model_dict.update(state_dict)model.load_state_dict(model_dict)

關(guān)于PyTorch模型的保存與加載是怎么樣的問題的解答就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，如果你還有很多疑惑沒有解開，可以關(guān)注創(chuàng)新互聯(lián)行業(yè)資訊頻道了解更多相關(guān)知識(shí)。

當(dāng)前標(biāo)題：PyTorch模型的保存與加載是怎么樣的
瀏覽地址：http://aaarwkj.com/article22/gooecc.html

成都網(wǎng)站建設(shè)公司_創(chuàng)新互聯(lián)，為您提供網(wǎng)站建設(shè)、建站公司、響應(yīng)式網(wǎng)站、Google、外貿(mào)網(wǎng)站建設(shè)、手機(jī)網(wǎng)站建設(shè)

廣告

聲明：本網(wǎng)站發(fā)布的內(nèi)容（圖片、視頻和文字）以用戶投稿、用戶轉(zhuǎn)載內(nèi)容為主，如果涉及侵權(quán)請(qǐng)盡快告知，我們將會(huì)在第一時(shí)間刪除。文章觀點(diǎn)不代表本網(wǎng)站立場(chǎng)，如需處理請(qǐng)聯(lián)系客服。電話：028-86922220；郵箱：631063699@qq.com。內(nèi)容未經(jīng)允許不得轉(zhuǎn)載，或轉(zhuǎn)載時(shí)需注明來源： 創(chuàng)新互聯(lián)