一、模型參數(shù)的保存和加載

• torch.save(module.state_dict(), path)：使用module.state_dict()函數(shù)獲取各層已經(jīng)訓(xùn)練好的參數(shù)和緩沖區(qū)，然后將參數(shù)和緩沖區(qū)保存到path所指定的文件存放路徑（常用文件格式為.pt、.pth或.pkl）。
• torch.nn.Module.load_state_dict(state_dict)：從state_dict中加載參數(shù)和緩沖區(qū)到Module及其子類中 。
• torch.nn.Module.state_dict()函數(shù)返回python中的一個OrderedDict類型字典對象，該對象將每一層與它的對應(yīng)參數(shù)和緩沖區(qū)建立映射關(guān)系，字典的鍵值是參數(shù)或緩沖區(qū)的名稱。只有那些參數(shù)可以訓(xùn)練的層才會被保存到OrderedDict中，例如：卷積層、線性層等。
• Python中的字典類以“鍵:值”方式存取數(shù)據(jù)，OrderedDict是它的一個子類，實現(xiàn)了對字典對象中元素的排序（OrderedDict根據(jù)放入元素的先后順序進行排序）。由于進行了排序，所以順序不同的兩個OrderedDict字典對象會被當(dāng)做是兩個不同的對象。
• 示例：

import torch
import torch.nn as nn
class Net(nn.Module):
 def __init__(self):
  super(Net, self).__init__()
  self.conv1 = nn.Conv2d(1, 2, 3)
  self.pool1 = nn.MaxPool2d(2, 2)
 def forward(self, x):
  x = self.conv1(x)
  x = self.pool1(x)
  return x
# 初始化網(wǎng)絡(luò)
net = Net()
net.conv1.weight[0].detach().fill_(1)
net.conv1.weight[1].detach().fill_(2)
net.conv1.bias.data.detach().zero_()
# 獲取state_dict
state_dict = net.state_dict()
# 字典的遍歷默認(rèn)是遍歷key，所以param_tensor實際上是鍵值
for param_tensor in state_dict: 
 print(param_tensor,'：\n',state_dict[param_tensor])
# 保存模型參數(shù)
torch.save(state_dict,"net_params.pth")
# 通過加載state_dict獲取模型參數(shù)
net.load_state_dict(state_dict)

輸出：

二、完整模型的保存和加載

• torch.save(module, path)：將訓(xùn)練完的整個網(wǎng)絡(luò)模型module保存到path所指定的文件存放路徑（常用文件格式為.pt或.pth）。
• torch.load(path)：加載保存到path中的整個神經(jīng)網(wǎng)絡(luò)模型。
• 示例：

import torch
import torch.nn as nn
class Net(nn.Module):
 def __init__(self):
  super(Net, self).__init__()
  self.conv1 = nn.Conv2d(1, 2, 3)
  self.pool1 = nn.MaxPool2d(2, 2)
 def forward(self, x):
  x = self.conv1(x)
  x = self.pool1(x)
  return x
# 初始化網(wǎng)絡(luò)
net = Net()
net.conv1.weight[0].detach().fill_(1)
net.conv1.weight[1].detach().fill_(2)
net.conv1.bias.data.detach().zero_()
# 保存整個網(wǎng)絡(luò)
torch.save(net,"net.pth")
# 加載網(wǎng)絡(luò)
net = torch.load("net.pth")

到此這篇關(guān)于PyTorch深度學(xué)習(xí)模型的保存和加載流程詳解的文章就介紹到這了,更多相關(guān)PyTorch 模型的保存 內(nèi)容請搜索本站以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持本站！

版權(quán)聲明：本站文章來源標(biāo)注為YINGSOO的內(nèi)容版權(quán)均為本站所有，歡迎引用、轉(zhuǎn)載，請保持原文完整并注明來源及原文鏈接。禁止復(fù)制或仿造本網(wǎng)站，禁止在非maisonbaluchon.cn所屬的服務(wù)器上建立鏡像，否則將依法追究法律責(zé)任。本站部分內(nèi)容來源于網(wǎng)友推薦、互聯(lián)網(wǎng)收集整理而來，僅供學(xué)習(xí)參考，不代表本站立場，如有內(nèi)容涉嫌侵權(quán)，請聯(lián)系alex-e#qq.com處理。