注意
按一下這裡以下載完整的範例程式碼
儲存和載入模型¶
建立於：2018 年 8 月 29 日 | 最後更新：2024 年 9 月 10 日 | 最後驗證：2024 年 11 月 05 日
本文檔提供了有關儲存和載入 PyTorch 模型的各種使用案例的解決方案。您可以自由閱讀整個文檔，或者直接跳到您需要的程式碼以獲得所需的使用案例。
在儲存和載入模型時，需要熟悉三個核心函數
torch.save：將序列化物件儲存到磁碟。此函數使用 Python 的 pickle 工具進行序列化。可以使用此函數儲存模型、張量和各種物件的字典。
torch.load：使用 pickle 的 unpickling 功能將 pickled 物件檔案反序列化到記憶體。此函數還有助於將資料載入到裝置中（請參閱 跨裝置儲存和載入模型）。
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 和批次標準化層設定為評估模式。否則會產生不一致的推論結果。
注意
請注意，load_state_dict() 函數採用字典物件，而不是已儲存物件的路徑。這表示您必須先反序列化已儲存的 state_dict，然後才能將其傳遞給 load_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 和批次標準化層設定為評估模式。否則會產生不一致的推論結果。
以 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 和批次標準化層設定為評估模式。否則會產生不一致的推論結果。
有關 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 和批次正規化層設定為評估模式。否則會產生不一致的推論結果。如果您希望恢復訓練，請呼叫 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 和批次正規化層設定為評估模式。否則會產生不一致的推論結果。如果您希望恢復訓練，請呼叫 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。 因此，請記得手動覆寫 tensors： 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'))，將模型的參數 tensors 轉換為 CUDA tensors。 最後，請務必對所有模型輸入使用 .to(torch.device('cuda')) 函數，以便為 CUDA 優化模型準備資料。 請注意，呼叫 my_tensor.to(device) 會在 GPU 上傳回 my_tensor 的新副本。 它不會覆寫 my_tensor。 因此，請記得手動覆寫 tensors： 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 秒）
