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

1. torch.save：將序列化對象保存到磁盤。此函數使用Python的pickle模塊進行序列化。使用此函數可以保存如模型、tensor、字典等各種對象。
2. torch.load：使用pickle的unpickling功能，將pickle對象文件反序列化到內存。此功能還可以有助于設備加載數據。
3. torch.nn.Module.load_state_dict：使用反序列化函數 state_dict 來加載模型的參數字典。
4. 什么是狀態字典：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()來加載本地字典。這里，可以非常容易的通過簡單查詢字典來訪問所保存的項目。
在運行推理之前，務必調用model.eval()去設置 dropout 和 batch normalization 為評估。如果不這樣做，有可能得到不一致的推斷結果。 如果想要恢復訓練，請調用model.train()以確保這些層處于訓練模式。
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()。 這樣，就可以非常靈活地以任何方式加載模型到想要的設備中。

總結

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

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