PyTorch框架學習十九——模型加載與保存

• 一、序列化與反序列化
• 二、PyTorch中的序列化與反序列化
• • 1.torch.save
  • 2.torch.load
• 三、模型的保存
• • 1.方法一：保存整個Module
  • 2.方法二：僅保存模型參數
• 四、模型的加載
• • 1.加載整個模型
  • 2.僅加載模型參數
• 五、斷點續訓練
• • 1.斷點續訓練的模型保存
  • 2.斷點續訓練的模型加載

距離上次的學習筆記時隔了正好一個月。。。下面繼續！

其實深度學習的五大步驟——數據、模型、損失函數、優化器和迭代訓練，這些用PyTorch的基本實現前面的筆記都已經涉及，最后的幾次筆記是一些實用的技巧，PyTorch框架的基礎學習就將告一段落，后面開始Paper的學習，有精力的話也一定會更新。

一、序列化與反序列化

這里先介紹一下內存與硬盤的區別，這兩種存儲設備最重要的區別是內存的內容不會永久保存，在程序運行結束或是斷電等情況下，數據會全部丟失，而在硬盤里的數據則不受上述的影響。所以我們有必要將一些值得長久保存的東西，比如訓練好的模型、關鍵的數據等從內存中轉存到硬盤中，下次使用的時候再從硬盤中讀取到內存。

• 序列化：將數據從內存轉存到硬盤。
• 反序列化：將數據從硬盤轉存到內存。

具體見下圖所示：

序列化與反序列化的主要目的是：可以使得數據/模型可以長久地保存。

二、PyTorch中的序列化與反序列化

1.torch.save

torch.save(obj, f, pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, pickle_protocol=2, _use_new_zipfile_serialization=True)

參數如下所示：

主要需要關注兩個：

• obj：需要保存的對象。
• f：保存的路徑。

2.torch.load

torch.load(f, map_location=None, pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, **pickle_load_args)

參數如下所示：

主要需要關注前兩個：

• f：加載文件的路徑。
• map_location：指定存放的位置，cpu或gpu。

三、模型的保存

1.方法一：保存整個Module

torch.save(net, path)

下面構建一個簡單的網絡模型作為例子：

import torch import numpy as np import torch.nn as nn# 構建一個LeNet網絡模型，并實例化為net class LeNet2(nn.Module):def __init__(self, classes):super(LeNet2, self).__init__()self.features = nn.Sequential(nn.Conv2d(3, 6, 5),nn.ReLU(),nn.MaxPool2d(2, 2),nn.Conv2d(6, 16, 5),nn.ReLU(),nn.MaxPool2d(2, 2))self.classifier = nn.Sequential(nn.Linear(16*5*5, 120),nn.ReLU(),nn.Linear(120, 84),nn.ReLU(),nn.Linear(84, classes))def forward(self, x):x = self.features(x)x = x.view(x.size()[0], -1)x = self.classifier(x)return xdef initialize(self):for p in self.parameters():p.data.fill_(20191104)net = LeNet2(classes=2019)

忽略訓練過程，真實的訓練不應該這樣。。。，這里為了簡單，只是人為地改變了一下參數權重，并不是訓練得到的結果。

# "訓練" print("訓練前: ", net.features[0].weight[0, ...]) net.initialize() print("訓練后: ", net.features[0].weight[0, ...])

（虛假）結果：

訓練前: tensor([[[-0.0065, -0.0388, 0.0194, -0.0944, -0.0539],[-0.0981, 0.0611, -0.0066, -0.0140, 0.0849],[ 0.0495, -0.0793, -0.0424, 0.0282, 0.0338],[ 0.0523, -0.0409, -0.1071, -0.0623, 0.0956],[-0.0385, 0.0554, 0.0160, 0.0671, 0.0913]],[[ 0.0409, 0.0811, -0.0677, 0.1040, -0.0236],[-0.0717, -0.0418, 0.0007, 0.0950, -0.0309],[-0.0089, 0.0140, -0.0855, 0.0818, -0.0270],[ 0.0316, -0.0637, 0.0093, 0.0669, -0.1031],[-0.0337, 0.1068, -0.0927, -0.1069, 0.0267]],[[-0.0034, 0.0862, 0.0743, -0.0082, 0.0929],[-0.0738, 0.0781, -0.0701, -0.0327, -0.0553],[-0.0951, -0.0604, -0.0906, 0.0169, 0.0084],[ 0.0589, 0.0391, 0.0493, 0.1018, 0.0563],[ 0.0407, -0.0053, 0.0307, 0.0077, 0.0262]]],grad_fn=<SelectBackward>) 訓練后: tensor([[[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.]],[[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.]],[[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.]]],grad_fn=<SelectBackward>)

好，通過上面的偷懶操作，視為我們訓練好了模型，現在需要保存這個模型，方法一是保存整個模型，如下圖所示：

path_model = "./model.pkl" torch.save(net, path_model)

此時查看當前目錄，多了一個model.pkl的文件，這個就是保存的模型文件：

這是一個簡單的模型，如果模型比較龐大的話，這種方法會有一個問題，就是保存耗時較長，占用內存也較大，所以一般不這樣保存。

2.方法二：僅保存模型參數

我們回顧一下模型的內容，一個模型它是有很多部分組成的，就像剛剛構建的LeNet，其模型所包含的內容就如下所示：

而我們可以觀察一下，在模型的訓練過程中，只有參數_parameters是不斷被更新的，那么我們可以在保存的時候僅保存模型參數，其他的部分在用的時候可以重新構建，然后將參數加載到新構建的模型上就等同于是保存了整個模型，這樣不僅可以節省保存的時間也能減小內存加載的消耗。

這種方法的保存代碼為：

net_state_dict = net.state_dict() torch.save(net_state_dict, path)

這里使用了state_dict()函數，返回了一個字典，包含了這個模型的所有參數。

還是以剛剛的LeNet為例，只需將保存的代碼改為：

path_state_dict = "./model_state_dict.pkl" # 保存模型參數 net_state_dict = net.state_dict() torch.save(net_state_dict, path_state_dict)

即可。

可以通過單步調試查看net_state_dict 的內容：

它的確包含了每一層的權重和偏置，而在當前目錄下多了一個名為model_state_dict.pkl的保存文件。

四、模型的加載

這個都是和三中的兩種方法對應，保存的時候是整個模型就加載整個模型，保存的時候是僅保存模型參數就先構建模型然后加載模型參數。

1.加載整個模型

import torch import numpy as np import torch.nn as nnclass LeNet2(nn.Module):def __init__(self, classes):super(LeNet2, self).__init__()self.features = nn.Sequential(nn.Conv2d(3, 6, 5),nn.ReLU(),nn.MaxPool2d(2, 2),nn.Conv2d(6, 16, 5),nn.ReLU(),nn.MaxPool2d(2, 2))self.classifier = nn.Sequential(nn.Linear(16*5*5, 120),nn.ReLU(),nn.Linear(120, 84),nn.ReLU(),nn.Linear(84, classes))def forward(self, x):x = self.features(x)x = x.view(x.size()[0], -1)x = self.classifier(x)return xdef initialize(self):for p in self.parameters():p.data.fill_(20191104)# ================================== load net =========================== flag = 1 # flag = 0 if flag:path_model = "./model.pkl"net_load = torch.load(path_model)print(net_load)

輸出：

LeNet2((features): Sequential((0): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))(1): ReLU()(2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(3): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))(4): ReLU()(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))(classifier): Sequential((0): Linear(in_features=400, out_features=120, bias=True)(1): ReLU()(2): Linear(in_features=120, out_features=84, bias=True)(3): ReLU()(4): Linear(in_features=84, out_features=2019, bias=True)) )

注意：在加載之前一定要先定義模型LeNet，定義必須和保存時的一致，不然會對應不上。

2.僅加載模型參數

這種方法需要比前一種方法多一步：實例化模型。即先得有一個模型才能對它加載參數。

# ================================== load state_dict ===========================flag = 1 # flag = 0 if flag:# 加載參數到內存，并查看keyspath_state_dict = "./model_state_dict.pkl"state_dict_load = torch.load(path_state_dict)print(state_dict_load.keys())# 實例化模型net_new = LeNet2(classes=2019)# 更新參數到模型，等同于加載整個模型print("加載前: ", net_new.features[0].weight[0, ...])net_new.load_state_dict(state_dict_load)print("加載后: ", net_new.features[0].weight[0, ...])

輸出：

odict_keys(['features.0.weight', 'features.0.bias', 'features.3.weight', 'features.3.bias', 'classifier.0.weight', 'classifier.0.bias', 'classifier.2.weight', 'classifier.2.bias', 'classifier.4.weight', 'classifier.4.bias']) 加載前: tensor([[[ 0.1004, -0.0658, -0.0483, -0.0398, 0.0197],[ 0.0157, 0.0474, -0.0996, -0.0778, -0.0177],[-0.0889, -0.0732, -0.0625, 0.1042, 0.1082],[ 0.0360, -0.0571, 0.0523, -0.0154, -0.0241],[ 0.0427, 0.0968, 0.0166, 0.0405, -0.0782]],[[ 0.0610, 0.0467, 0.0285, 0.0061, 0.0146],[-0.0410, 0.0822, -0.0535, 0.0842, -0.0443],[ 0.1010, 0.0972, -0.0625, -0.1114, 0.0899],[ 0.0462, -0.0121, 0.0476, 0.0673, 0.0900],[ 0.1030, -0.0278, 0.0374, -0.0098, -0.0929]],[[ 0.0703, 0.0600, -0.0538, 0.0179, -0.0008],[ 0.1119, 0.0187, 0.0359, -0.0721, 0.0033],[-0.0225, 0.0788, -0.1081, 0.0700, 0.0269],[ 0.0519, -0.0959, -0.0329, -0.0520, -0.0741],[ 0.1020, 0.0321, 0.0549, -0.0522, -0.0783]]],grad_fn=<SelectBackward>) 加載后: tensor([[[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.]],[[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.]],[[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.],[20191104., 20191104., 20191104., 20191104., 20191104.]]],grad_fn=<SelectBackward>)

五、斷點續訓練

這是一個非常實用的技巧，設想你的一個模型訓練了五天，然后因為斷電等不可控原因突然中止，是不是非常抓狂。

斷點續訓練的思想就是，在訓練的過程中，每隔一定數量的epoch就保存一次模型（參數），若程序中斷，可以從離中斷最近位置保存的文件中加載模型（參數）繼續訓練。

1.斷點續訓練的模型保存

在訓練過程中，只有模型與優化器中會有與訓練相關的數據，比如模型的參數權重、優化器中的學習率等等，以及還需要記錄一下保存時的epoch，所以我們可以構建這樣一個字典，存放剛剛所說的三類數據：

checkpoint = {"model_state_dict": net.state_dict(),"optimizer_state_dict": optimizer.state_dict(),"epoch": epoch}

將這個checkpoint保存為pkl文件即可。

if (epoch+1) % checkpoint_interval == 0:checkpoint = {"model_state_dict": net.state_dict(),"optimizer_state_dict": optimizer.state_dict(),"epoch": epoch}path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)torch.save(checkpoint, path_checkpoint)

其中checkpoint_interval 就是每隔多少個epoch保存一次上述信息。

下面我們人為構建一個中斷，模型、損失函數、優化器、迭代訓練部分的代碼省略：

if epoch > 5:print("訓練意外中斷...")break

結果如下所示：

Training:Epoch[000/010] Iteration[010/012] Loss: 0.6753 Acc:53.75% Valid: Epoch[000/010] Iteration[003/003] Loss: 0.4504 Acc:57.37% Training:Epoch[001/010] Iteration[010/012] Loss: 0.3928 Acc:83.75% Valid: Epoch[001/010] Iteration[003/003] Loss: 0.0130 Acc:85.79% Training:Epoch[002/010] Iteration[010/012] Loss: 0.1717 Acc:93.12% Valid: Epoch[002/010] Iteration[003/003] Loss: 0.0027 Acc:92.11% Training:Epoch[003/010] Iteration[010/012] Loss: 0.1100 Acc:95.62% Valid: Epoch[003/010] Iteration[003/003] Loss: 0.0003 Acc:96.32% Training:Epoch[004/010] Iteration[010/012] Loss: 0.0321 Acc:98.12% Valid: Epoch[004/010] Iteration[003/003] Loss: 0.0000 Acc:98.42% Training:Epoch[005/010] Iteration[010/012] Loss: 0.0064 Acc:100.00% Valid: Epoch[005/010] Iteration[003/003] Loss: 0.0000 Acc:100.00% Training:Epoch[006/010] Iteration[010/012] Loss: 0.0109 Acc:99.38% 訓練意外中斷...

當前目錄下多了一個pkl文件，其中的4是指保存的第五輪訓練后的結果：

2.斷點續訓練的模型加載

與正常訓練差不多，只是需要在迭代訓練之前把上述保存的三種信息加載給模型，使之恢復到斷點的狀態：

path_checkpoint = "./checkpoint_4_epoch.pkl" checkpoint = torch.load(path_checkpoint) # 模型參數更新 net.load_state_dict(checkpoint['model_state_dict']) # 優化器參數更新 optimizer.load_state_dict(checkpoint['optimizer_state_dict']) # epoch數恢復到斷點 start_epoch = checkpoint['epoch'] scheduler.last_epoch = start_epoch

結果如下所示：

Training:Epoch[005/010] Iteration[010/012] Loss: 0.0146 Acc:99.38% Valid: Epoch[005/010] Iteration[003/003] Loss: 0.0000 Acc:99.47% Training:Epoch[006/010] Iteration[010/012] Loss: 0.0039 Acc:100.00% Valid: Epoch[006/010] Iteration[003/003] Loss: 0.0000 Acc:100.00% Training:Epoch[007/010] Iteration[010/012] Loss: 0.0013 Acc:100.00% Valid: Epoch[007/010] Iteration[003/003] Loss: 0.0000 Acc:100.00% Training:Epoch[008/010] Iteration[010/012] Loss: 0.0020 Acc:100.00% Valid: Epoch[008/010] Iteration[003/003] Loss: 0.0000 Acc:100.00% Training:Epoch[009/010] Iteration[010/012] Loss: 0.0010 Acc:100.00% Valid: Epoch[009/010] Iteration[003/003] Loss: 0.0000 Acc:100.00%

可見，epoch繼續從5開始，準確率一開始就很高，是緊接著前5次的結果的，再觀察損失函數的變化，初始值就很小了：

綜上，這樣的方法，可以使得因特殊原因中途中斷的訓練從中斷位置繼續訓練而不用從頭開始，非常有用。

總結

以上是生活随笔為你收集整理的PyTorch框架学习十九——模型加载与保存的全部內容，希望文章能夠幫你解決所遇到的問題。

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