眾所周知，python的對象都可以通過torch.save和torch.load函數進行保存和加載(不知道？那你現在知道了(*^_^*))，比如：

x1 = {"d":"ddf","dd":'fdsf'} torch.save(x1, 'a1.pt')x2 = ["ddf",'fdsf'] torch.save(x2, 'a2.pt')x3 = 1 torch.save(x3, 'a3.pt')x4 = torch.ones(3) torch.save(x4, 'a4.pt')

讀取的時候也是一樣：

x5 = torch.load('a1.pt')x6 = torch.load('a2.pt')x7 = torch.load('a3.pt')x8 = torch.load('a4.pt')

這種非常簡單粗暴，直接把整個對象扔進磁盤文件里保存，所以對于我們訓練好的模型來說，因為訓練好的模型也是一個對象，所以我們也可以使用這個方法把訓練好的模型對象直接扔進去。但是這樣有一個問題，就是模型對象開銷比較大，比如最近包含1350億個參數的那個有名的神經網絡模型，如果把它保存到磁盤里面沒有百八十T是保存不下的。所以我們是不是可以僅僅保存模型里面的關鍵數據呢？

答案是，可以！

因為決定一個模型是什么樣有兩方面的因素，一個是模型的結構是什么，另一個是模型的參數是什么，這兩個定了，這個模型也就確定了。模型的結構在我們初始化模型對象的時候就定了，比如對于任意一個模型類，我們初始化它的兩個對象，這兩個對象代表的模型的結構肯定是一樣的，區別就在于它們的參數不一樣。所以我們保存模型的關鍵就是保存模型的參數，而模型的結構每次用的時候新建一個對象就好了，然后從磁盤里把模型的參數讀取出來賦給這個對象。是不是超級簡單？

那我們怎么拿到模型的參數呢？巧了！

模型的state_dict()函數就是返回模型的所有參數的（這個函數是nn.Module的，所以所有繼承了nn.Module的模型類都有這個函數），比如：

class MLP(nn.Module):def __init__(self):super(MLP, self).__init__()self.hidden = nn.Linear(3, 2)self.act = nn.ReLU()self.output = nn.Linear(2, 1)def forward(self, x):a = self.act(self.hidden(x))return self.output(a)net = MLP() net.state_dict()

輸出：

OrderedDict([('hidden.weight',tensor([[-0.4195, 0.2609, 0.4325],[-0.4031, 0.2078, 0.2077]])),('hidden.bias', tensor([ 0.0755, -0.1408])),('output.weight', tensor([[0.2473, 0.6614]])),('output.bias', tensor([0.6191]))])

有的同學可能注意到了，self.act層的參數沒有包含進來！

大哥，self.act層沒有參數好嗎（捂臉）

還有的同學可能想問，那有的層有參數、有的層沒有參數，那萬一加載的時候把某個參數給錯了層怎么辦？

完全不會！注意看，state_dict()返回的是一個字典，每一個張量都對應的有層的名字，清清楚楚，絕對沒有問題。

那這樣就簡單了，舉個例子看一下：

X = torch.randn(2, 3) Y = net(X) # 這個net就是上面創建的那個對象，我們把它的參數保存起來，然后新建一個net2，然后把保存的這些參數加載進net2，這樣我們把X輸入net2得到的Y2應該與Y是相等的PATH = "./net.pt" torch.save(net.state_dict(), PATH)net2 = MLP() net2.load_state_dict(torch.load(PATH)) Y2 = net2(X) Y2 == Y

輸出：

tensor([[1],[1]], dtype=torch.uint8)

輸出的張量，代表Y2 == Y比較結果為true，也就是說是一樣的，驗證了我們的猜想（上面代碼注釋中的那個猜想）。

好了，以上就是pytorch保存和加載模型的兩種方法，是不是非常簡單？

陽陽：保存和加載pytorch模型的兩種方法，選哪個好？?zhuanlan.zhihu.com

總結

以上是生活随笔為你收集整理的加载dict_PyTorch 7.保存和加载pytorch模型的两种方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。