開始使用免費開始

儲存與載入模型

一家製造公司想根據影像對其專案進行分類,並決定合適的出貨包裝。你已在 PyTorch 訓練出高準確率的模型,現在打算將模型與其預先訓練的權重儲存起來,便於未來使用並與團隊分享,確保他們能順利載入。

torchtorch.nn(作為 nn)已匯入。預先訓練的模型物件在你的工作空間中名為 model,其架構為 ManufacturingCNN

本練習屬於課程

使用 PyTorch 進行影像深度學習

檢視課程

練習說明

  • 將預先訓練的模型儲存為 ModelCNN.pth,記得要儲存權重而不只是架構。
  • ManufacturingCNN() 類別建立名為 loaded_model 的模型實例。
  • ModelCNN.pth 的權重載入到 loaded_model,並將權重傳入 .load_state_dict()

動手互動練習

試著完成這個範例程式碼,體驗一下這個練習。

# Save the model
torch.____(model.____, ____)

# Create a new model
loaded_model = ____

# Load the saved model
loaded_model.____(torch.____('ModelCNN.pth'))
print(loaded_model)
編輯並執行程式碼