モデルの保存と読み込み
製造業の会社が、画像に基づいてプロジェクトを分類し、適切な出荷用の梱包を決めようとしています。PyTorch で高精度なモデルを学習できたので、将来の再利用やチームへの共有に向けて、モデル本体と事前学習済みの重みを保存し、だれでもシームレスに読み込めるようにしたいと考えています。
torch と torch.nn(nn)はインポート済みです。事前学習済みのモデルオブジェクトは model、そのアーキテクチャは ManufacturingCNN としてワークスペースに用意されています。
この演習はコースの一部です
PyTorch で学ぶ画像向け Deep Learning
演習の手順
- 事前学習済みモデルを
ModelCNN.pthとして保存します。アーキテクチャだけでなく、必ず重みも保存してください。 - クラス
ManufacturingCNN()からモデルインスタンスloaded_modelを作成します。 .load_state_dict()に重みを渡して、ModelCNN.pthの重みをloaded_modelに読み込みます。
実践的なインタラクティブ演習
このサンプルコードを完成させて、この演習に挑戦してみましょう。
# Save the model
torch.____(model.____, ____)
# Create a new model
loaded_model = ____
# Load the saved model
loaded_model.____(torch.____('ModelCNN.pth'))
print(loaded_model)