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