在 PyTorch 中加載用於推理的遷移學習模型的正確方法是什么？

Question

我正在使用基於 Resnet152 的遷移學習來訓練模型。 基於 PyTorch 教程，我在保存經過訓練的模型並加載它以進行推理方面沒有問題。 但是，加載模型所需的時間很慢。 我不知道我做對了沒有，這是我的代碼：

將訓練好的模型保存為狀態字典：

torch.save(model.state_dict(), 'model.pkl')

加載它以進行推理：

model = models.resnet152()
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(classes))
st = torch.load('model.pkl', map_location='cuda:0' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(st)
model.eval()

我對代碼進行了計時，發現第一行model = models.resnet152()加載時間最長。 在 CPU 上，測試一張圖像需要 10 秒。 所以我的想法是這可能不是加載它的正確方法？

如果我像這樣保存整個模型而不是 state.dict：

torch.save(model, 'model_entire.pkl')

並像這樣測試它：

model = torch.load('model_entire.pkl')
model.eval()

在同一台機器上，測試一張圖像只需 5 秒。

所以我的問題是：這是加載 state_dict 進行推理的正確方法嗎？

Answer 1

在第一個代碼片段中，您從 TorchVision 下載模型（具有隨機權重），然后將您的（本地存儲的）權重加載到它。

在第二個示例中，您正在加載本地存儲的模型（及其權重）。

前者會更慢，因為您需要連接到托管模型的服務器並下載它，而不是本地文件，但它更易於復制，不依賴於您的本地文件。 此外，時間差應該是一次性初始化，並且它們應該具有相同的時間復雜度（因為在您執行推理時，模型已經加載到兩者中，並且它們是等效的）。