PyTorch專欄（七）:模型保存與載入那些事

2019 年 10 月 4 日
筆記

作者 | News

編輯 | 安可

出品 | 磐創AI團隊出品

【磐創AI 導讀】：本篇文章講解了PyTorch專欄的第三章中的保存和載入模型。查看專欄歷史文章，請點擊下方藍色字體進入相應鏈接閱讀。查看關於本專欄的介紹：PyTorch專欄開篇。

專欄目錄：

第一章：PyTorch之簡介與下載

PyTorch簡介
PyTorch環境搭建

第二章：PyTorch之60分鐘入門

PyTorch入門
PyTorch自動微分
PyTorch神經網路
PyTorch影像分類器
PyTorch數據並行處理

第三章：PyTorch之入門強化

第四章：PyTorch之影像篇

微調基於torchvision 0.3的目標檢測模型
微調TorchVision模型
空間變換器網路
使用PyTorch進行神經傳遞
生成對抗示例
使用ONNX將模型轉移至Caffe2和移動端

第五章：PyTorch之文本篇

聊天機器人教程
使用字元級RNN生成名字
使用字元級RNN進行名字分類
在深度學習和NLP中使用Pytorch
使用Sequence2Sequence網路和注意力進行翻譯

第六章：PyTorch之生成對抗網路

第七章：PyTorch之強化學習

當保存和載入模型時，需要熟悉三個核心功能：

torch.save：將序列化對象保存到磁碟。此函數使用Python的pickle模組進行序列化。使用此函數可以保存如模型、tensor、字典等各種對象。
torch.load：使用pickle的unpickling功能將pickle對象文件反序列化到記憶體。此功能還可以有助於設備載入數據。
torch.nn.Module.load_state_dict：使用反序列化函數 state_dict 來載入模型的參數字典。

1.什麼是狀態字典：state_dict?

在PyTorch中，torch.nn.Module模型的可學習參數（即權重和偏差）包含在模型的參數中，（使用model.parameters()可以進行訪問）。 state_dict是Python字典對象，它將每一層映射到其參數張量。注意，只有具有可學習參數的層（如卷積層，線性層等）的模型才具有state_dict這一項。目標優化torch.optim也有state_dict屬性，它包含有關優化器的狀態資訊，以及使用的超參數。

因為state_dict的對象是Python字典，所以它們可以很容易的保存、更新、修改和恢復，為PyTorch模型和優化器添加了大量模組。

下面通過從簡單模型訓練一個分類器中來了解一下state_dict的使用。

# 定義模型  class TheModelClass(nn.Module):      def __init__(self):          super(TheModelClass, self).__init__()          self.conv1 = nn.Conv2d(3, 6, 5)          self.pool = nn.MaxPool2d(2, 2)          self.conv2 = nn.Conv2d(6, 16, 5)          self.fc1 = nn.Linear(16 * 5 * 5, 120)          self.fc2 = nn.Linear(120, 84)          self.fc3 = nn.Linear(84, 10)        def forward(self, x):          x = self.pool(F.relu(self.conv1(x)))          x = self.pool(F.relu(self.conv2(x)))          x = x.view(-1, 16 * 5 * 5)          x = F.relu(self.fc1(x))          x = F.relu(self.fc2(x))          x = self.fc3(x)          return x    # 初始化模型  model = TheModelClass()    # 初始化優化器  optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)    # 列印模型的狀態字典  print("Model's state_dict:")  for param_tensor in model.state_dict():      print(param_tensor, "t", model.state_dict()[param_tensor].size())    # 列印優化器的狀態字典  print("Optimizer's state_dict:")  for var_name in optimizer.state_dict():      print(var_name, "t", optimizer.state_dict()[var_name])

輸出

Model's state_dict:  conv1.weight     torch.Size([6, 3, 5, 5])  conv1.bias   torch.Size([6])  conv2.weight     torch.Size([16, 6, 5, 5])  conv2.bias   torch.Size([16])  fc1.weight   torch.Size([120, 400])  fc1.bias     torch.Size([120])  fc2.weight   torch.Size([84, 120])  fc2.bias     torch.Size([84])  fc3.weight   torch.Size([10, 84])  fc3.bias     torch.Size([10])    Optimizer's state_dict:  state    {}  param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

2.保存和載入推理模型

2.1 保存/載入`state_dict`（推薦使用）

保存

torch.save(model.state_dict(), PATH)

載入

model = TheModelClass(*args, **kwargs)  model.load_state_dict(torch.load(PATH))  model.eval()

當保存好模型用來推斷的時候，只需要保存模型學習到的參數，使用torch.save()函數來保存模型state_dict,它會給模型恢復提供最大的靈活性，這就是為什麼要推薦它來保存的原因。

在 PyTorch 中最常見的模型保存使『.pt』或者是『.pth』作為模型文件擴展名。

請記住，在運行推理之前，務必調用model.eval()去設置 dropout 和 batch normalization 層為評估模式。如果不這麼做，可能導致模型推斷結果不一致。

注意

load_state_dict()函數只接受字典對象，而不是保存對象的路徑。這就意味著在你傳給load_state_dict()函數之前，你必須反序列化你保存的state_dict。例如，你無法通過 model.load_state_dict(PATH)來載入模型。

2.2 保存/載入完整模型

保存

torch.save(model, PATH)

載入

# 模型類必須在此之前被定義  model = torch.load(PATH)  model.eval()

此部分保存/載入過程使用最直觀的語法並涉及最少量的程式碼。以 Python `pickle 模組的方式來保存模型。這種方法的缺點是序列化數據受限於某種特殊的類而且需要確切的字典結構。這是因為pickle無法保存模型類本身。相反，它保存包含類的文件的路徑，該文件在載入時使用。因此，當在其他項目使用或者重構之後，您的程式碼可能會以各種方式中斷。

在 PyTorch 中最常見的模型保存使用『.pt』或者是『.pth』作為模型文件擴展名。

請記住，在運行推理之前，務必調用model.eval()設置 dropout 和 batch normalization 層為評估模式。如果不這麼做，可能導致模型推斷結果不一致。

3. 保存和載入 Checkpoint 用於推理/繼續訓練

保存

torch.save({              'epoch': epoch,              'model_state_dict': model.state_dict(),              'optimizer_state_dict': optimizer.state_dict(),              'loss': loss,              ...              }, PATH)

載入

model = TheModelClass(*args, **kwargs)  optimizer = TheOptimizerClass(*args, **kwargs)    checkpoint = torch.load(PATH)  model.load_state_dict(checkpoint['model_state_dict'])  optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  epoch = checkpoint['epoch']  loss = checkpoint['loss']    model.eval()  # - or -  model.train()

當保存成 Checkpoint 的時候，可用於推理或者是繼續訓練，保存的不僅僅是模型的 state_dict 。保存優化器的 state_dict 也很重要, 因為它包含作為模型訓練更新的緩衝區和參數。你也許想保存其他項目，比如最新記錄的訓練損失，外部的torch.nn.Embedding層等等。

要保存多個組件，請在字典中組織它們並使用torch.save()來序列化字典。PyTorch 中常見的保存checkpoint 是使用 .tar 文件擴展名。

要載入項目，首先需要初始化模型和優化器，然後使用torch.load()來載入本地字典。這裡,你可以非常容易的通過簡單查詢字典來訪問你所保存的項目。

請記住在運行推理之前，務必調用model.eval()去設置 dropout 和 batch normalization 為評估。如果不這樣做，有可能得到不一致的推斷結果。如果你想要恢復訓練，請調用model.train()以確保這些層處於訓練模式。

4. 在一個文件中保存多個模型

保存

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 = TheModelBClass(*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()

當保存一個模型由多個torch.nn.Modules組成時，例如GAN(對抗生成網路)、sequence-to-sequence (序列到序列模型), 或者是多個模型融合, 可以採用與保存常規檢查點相同的方法。換句話說，保存每個模型的 state_dict 的字典和相對應的優化器。如前所述，可以通過簡單地將它們附加到字典的方式來保存任何其他項目，這樣有助於恢復訓練。

PyTorch 中常見的保存 checkpoint 是使用 .tar 文件擴展名。

要載入項目，首先需要初始化模型和優化器，然後使用torch.load()來載入本地字典。這裡，你可以非常容易的通過簡單查詢字典來訪問你所保存的項目。

5. 使用在不同模型參數下的熱啟動模式

保存

torch.save(modelA.state_dict(), PATH)

載入

modelB = TheModelBClass(*args, **kwargs)  modelB.load_state_dict(torch.load(PATH), strict=False)

在遷移學習或訓練新的複雜模型時，部分載入模型或載入部分模型是常見的情況。利用訓練好的參數，有助於熱啟動訓練過程，並希望幫助你的模型比從頭開始訓練能夠更快地收斂。

無論是從缺少某些鍵的 state_dict 載入還是從鍵的數目多於載入模型的 state_dict , 都可以通過在load_state_dict()函數中將strict參數設置為 False 來忽略非匹配鍵的函數。

如果要將參數從一個層載入到另一個層，但是某些鍵不匹配，主要修改正在載入的 state_dict 中的參數鍵的名稱以匹配要在載入到模型中的鍵即可。

6. 通過設備保存/載入模型

6.1 保存到 CPU、載入到 CPU

保存

torch.save(model.state_dict(), PATH)

載入

device = torch.device('cpu')  model = TheModelClass(*args, **kwargs)  model.load_state_dict(torch.load(PATH, map_location=device))

當從CPU上載入模型在GPU上訓練時, 將torch.device('cpu')傳遞給torch.load()函數中的map_location參數.在這種情況下，使用 map_location參數將張量下的存儲器動態的重新映射到CPU設備。

6.2 保存到 GPU、載入到 GPU

保存

torch.save(model.state_dict(), PATH)

載入

device = torch.device("cuda")  model = TheModelClass(*args, **kwargs)  model.load_state_dict(torch.load(PATH))  model.to(device)  # 確保在你提供給模型的任何輸入張量上調用input = input.to(device)

當在GPU上訓練並把模型保存在GPU，只需要使用model.to(torch.device('cuda'))，將初始化的 model 轉換為 CUDA 優化模型。另外，請務必在所有模型輸入上使用.to(torch.device('cuda'))函數來為模型準備數據。請注意，調用my_tensor.to(device)會在GPU上返回my_tensor的副本。因此，請記住手動覆蓋張量：my_tensor= my_tensor.to(torch.device('cuda'))。

6.3 保存到 CPU，載入到 GPU

保存

torch.save(model.state_dict(), PATH)

載入

device = torch.device("cuda")  model = TheModelClass(*args, **kwargs)  model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want  model.to(device)  # 確保在你提供給模型的任何輸入張量上調用input = input.to(device)

在CPU上訓練好並保存的模型載入到GPU時，將torch.load()函數中的map_location參數設置為cuda:device_id。這會將模型載入到指定的GPU設備。接下來，請務必調用model.to(torch.device('cuda'))將模型的參數張量轉換為 CUDA 張量。最後，確保在所有模型輸入上使用 .to(torch.device('cuda'))函數來為CUDA優化模型。請注意，調用my_tensor.to(device)會在GPU上返回my_tensor的新副本。它不會覆蓋my_tensor。因此，請手動覆蓋張量my_tensor = my_tensor.to(torch.device('cuda'))。

6.4 保存 `torch.nn.DataParallel` 模型

保存

torch.save(model.module.state_dict(), PATH)

載入

# 載入任何你想要的設備

torch.nn.DataParallel是一個模型封裝，支援並行GPU使用。要普通保存 DataParallel 模型, 請保存model.module.state_dict()。這樣，你就可以非常靈活地以任何方式載入模型到你想要的設備中。

Tags: