pytorchによってトレーニングされたネットワークモデルを保存して呼び出す方法

問題の説明

ディープニューラルネットワークモデルはトレーニングが非常に難しいため、トレーニング済みのpytorchネットワークモデルを保存し、次に使用するときに直接呼び出すことができます。このモデルはどのように保存する必要がありますか？

解決

Pytorchは主に、モデルパラメータを保存する方法とモデル全体を保存する方法の2つの方法を提供します。

方法1：モデルパラメータのみを保存する

#保存
torch.save(the_model.state_dict(), PATH)
#读取
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

方法2：モデル全体を保存する

#保存
torch.save(the_model, PATH)
#读取
the_model = torch.load(PATH)

保存されたファイルの推奨ファイル名は×××.tar形式であることに注意してください。たとえば、PATHの形式は「./model_file_name/the_model_name.tar」に設定できます。

 

この記事も参照できます。

https://blog.csdn.net/weixin_41680653/article/details/93768559