注意
點選此處下載完整示例程式碼
儲存和載入模型¶
建立日期：2018 年 8 月 29 日 | 最後更新：2024 年 9 月 10 日 | 最後驗證：2024 年 11 月 5 日
本文件提供了關於儲存和載入 PyTorch 模型的各種用例的解決方案。您可以通讀全文，或者直接跳到特定用例所需的程式碼部分。
在儲存和載入模型時，需要熟悉三個核心函式：
torch.save：將序列化物件儲存到磁碟。此函式使用 Python 的 pickle 工具進行序列化。可以使用此函式儲存模型、張量以及包含各種物件的字典。
torch.load：使用 pickle 的反序列化工具將 pickle 的物件檔案反序列化到記憶體。此函式也方便將資料載入到指定的裝置（參見 跨裝置儲存和載入模型）。
torch.nn.Module.load_state_dict：使用反序列化的 state_dict 載入模型的引數字典。有關 state_dict 的更多資訊，請參見 什麼是 state_dict？。
目錄
什麼是 state_dict？¶
在 PyTorch 中，torch.nn.Module 模型的可學習引數（即權重和偏差）包含在模型的 parameters 中（透過 model.parameters() 訪問）。state_dict 只是一個 Python 字典物件，它將每一層對映到其引數張量。請注意，只有具有可學習引數的層（卷積層、線性層等）和註冊的緩衝區（例如 batchnorm 的 running_mean）才會在模型的 state_dict 中有條目。最佳化器物件（torch.optim）也有一個 state_dict，其中包含關於最佳化器狀態以及使用的超引數的資訊。
由於 state_dict 物件是 Python 字典，因此可以輕鬆儲存、更新、修改和恢復，這為 PyTorch 模型和最佳化器增加了很大的模組化靈活性。
示例：¶
讓我們看看在 訓練分類器 教程中使用的簡單模型的 state_dict。
# Define model
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
# Initialize model
model = TheModelClass()
# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# Print model's state_dict
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
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]}]
儲存和載入模型用於推理¶
儲存/載入 state_dict（推薦）¶
儲存
torch.save(model.state_dict(), PATH)
載入
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, weights_only=True))
model.eval()
注意
PyTorch 1.6 版本將 torch.save 切換為使用新的基於 zip 檔案格式。 torch.load 仍然保留載入舊格式檔案的能力。如果出於任何原因想讓 torch.save 使用舊格式，請傳遞 kwarg 引數 _use_new_zipfile_serialization=False。
儲存模型用於推理時，只需要儲存訓練好的模型的學習引數。使用 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) 進行載入。
注意
如果您只打算保留表現最好的模型（根據獲得的驗證損失），請不要忘記 best_model_state = model.state_dict() 返回的是狀態的引用，而不是副本！您必須序列化 best_model_state 或使用 best_model_state = deepcopy(model.state_dict())，否則您最佳的 best_model_state 將在隨後的訓練迭代中不斷更新。結果，最終的模型狀態將是過擬合模型的狀態。
儲存/載入整個模型¶
儲存
torch.save(model, PATH)
載入
# Model class must be defined somewhere
model = torch.load(PATH, weights_only=False)
model.eval()
這種儲存/載入過程使用最直觀的語法，並且涉及的程式碼量最少。以這種方式儲存模型將使用 Python 的 pickle 模組儲存整個模組。這種方法的缺點是序列化的資料繫結到儲存模型時使用的特定類和確切的目錄結構。原因是 pickle 不儲存模型類本身。相反，它儲存包含該類的檔案的路徑，該路徑在載入時使用。因此，當在其他專案中使用或重構後，您的程式碼可能會以各種方式中斷。
一個常見的 PyTorch 約定是使用 .pt 或 .pth 副檔名來儲存模型。
請記住，在執行推理之前，您必須呼叫 model.eval() 將 dropout 和 batch normalization 層設定為評估模式。如果不這樣做，將導致推理結果不一致。
以 TorchScript 格式匯出/載入模型¶
使用訓練好的模型進行推理的一種常用方法是使用 TorchScript，它是 PyTorch 模型的一種中間表示，可以在 Python 中執行，也可以在 C++ 等高效能環境中執行。TorchScript 實際上是推薦用於規模化推理和部署的模型格式。
注意
使用 TorchScript 格式，您可以載入匯出的模型並執行推理，而無需定義模型類。
匯出
model_scripted = torch.jit.script(model) # Export to TorchScript
model_scripted.save('model_scripted.pt') # Save
載入
model = torch.jit.load('model_scripted.pt')
model.eval()
請記住，在執行推理之前，您必須呼叫 model.eval() 將 dropout 和 batch normalization 層設定為評估模式。如果不這樣做，將導致推理結果不一致。
有關 TorchScript 的更多資訊，請訪問專門的教程。您將熟悉跟蹤轉換，並學習如何在 C++ 環境中執行 TorchScript 模組。
儲存和載入用於推理和/或恢復訓練的通用檢查點¶
儲存：¶
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, weights_only=True)
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()
儲存用於推理和/或恢復訓練的通用檢查點時，除了模型的 state_dict 之外，還必須儲存更多內容。儲存最佳化器的 state_dict 也很重要，因為它包含模型訓練時更新的緩衝區和引數。您可能希望儲存的其他項包括您停止時的 epoch、最新記錄的訓練損失、外部 torch.nn.Embedding 層等。因此，此類檢查點通常比單獨的模型大 2~3 倍。
要儲存多個元件，請將它們組織在一個字典中，然後使用 torch.save() 序列化該字典。一個常見的 PyTorch 約定是使用 .tar 副檔名來儲存這些檢查點。
要載入這些項，首先初始化模型和最佳化器，然後使用 torch.load() 在本地載入字典。之後，您可以像預期那樣透過簡單查詢字典來輕鬆訪問已儲存的項。
請記住，在執行推理之前，您必須呼叫 model.eval() 將 dropout 和 batch normalization 層設定為評估模式。如果不這樣做，將導致推理結果不一致。如果您希望恢復訓練，請呼叫 model.train() 以確保這些層處於訓練模式。
將多個模型儲存到一個檔案¶
儲存：¶
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, weights_only=True)
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、序列到序列模型或模型整合，您遵循與儲存通用檢查點時相同的方法。換句話說，儲存一個字典，其中包含每個模型的 state_dict 和相應的最佳化器。如前所述，您可以儲存任何可能幫助您恢復訓練的其他項，只需將它們附加到字典中即可。
一個常見的 PyTorch 約定是使用 .tar 副檔名來儲存這些檢查點。
要載入這些模型，首先初始化模型和最佳化器，然後使用 torch.load() 在本地載入字典。之後，您可以像預期那樣透過簡單查詢字典來輕鬆訪問已儲存的項。
請記住，在執行推理之前，您必須呼叫 model.eval() 將 dropout 和 batch normalization 層設定為評估模式。如果不這樣做，將導致推理結果不一致。如果您希望恢復訓練，請呼叫 model.train() 將這些層設定為訓練模式。
使用不同模型的引數來熱啟動模型¶
儲存：¶
torch.save(modelA.state_dict(), PATH)
載入：¶
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH, weights_only=True), strict=False)
部分載入模型或載入部分模型是遷移學習或訓練新的複雜模型時的常見場景。利用訓練好的引數，即使只有少量可用，也將有助於熱啟動訓練過程，並有望幫助您的模型比從頭訓練更快地收斂。
無論您是從缺少某些鍵的部分 state_dict 載入，還是載入的 state_dict 包含比您要載入的模型更多的鍵，您都可以在 load_state_dict() 函式中將 strict 引數設定為 False，以忽略不匹配的鍵。
如果您想將引數從一層載入到另一層，但某些鍵不匹配，只需更改您要載入的 state_dict 中引數鍵的名稱，使其與您要載入到的模型中的鍵匹配即可。
跨裝置儲存和載入模型¶
在 GPU 上儲存，在 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, weights_only=True))
在 CPU 上載入使用 GPU 訓練的模型時，將 torch.device('cpu') 傳遞給 torch.load() 函式中的 map_location 引數。在這種情況下，張量底層的儲存使用 map_location 引數動態地重新對映到 CPU 裝置。
在 GPU 上儲存，在 GPU 上載入¶
儲存
torch.save(model.state_dict(), PATH)
載入
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, weights_only=True))
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model
在 GPU 上載入使用 GPU 訓練和儲存的模型時，只需使用 model.to(torch.device('cuda')) 將初始化的 model 轉換為 CUDA 最佳化模型。此外，請確保對所有模型輸入使用 .to(torch.device('cuda')) 函式，以便為模型準備資料。請注意，呼叫 my_tensor.to(device) 會在 GPU 上返回 my_tensor 的新副本。它不會覆蓋 my_tensor。因此，請記住手動覆蓋張量：my_tensor = my_tensor.to(torch.device('cuda'))。
在 CPU 上儲存，在 GPU 上載入¶
儲存
torch.save(model.state_dict(), PATH)
載入
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, weights_only=True, map_location="cuda:0")) # Choose whatever GPU device number you want
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model
在 GPU 上載入使用 CPU 訓練和儲存的模型時，將 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'))。
儲存 torch.nn.DataParallel 模型¶
儲存
torch.save(model.module.state_dict(), PATH)
載入
# Load to whatever device you want
torch.nn.DataParallel 是一個模型包裝器，用於實現 GPU 平行計算。要通用地儲存 DataParallel 模型，請儲存 model.module.state_dict()。這樣，您可以靈活地將模型載入到任何裝置上，並以任何您想要的方式載入。
指令碼總執行時間： ( 0 分鐘 0.000 秒)
