以前にJetson Xavier NXとDockerでpytorch画像分類フレームワークのclassy visionの環境構築を行いました
technoxs-stacker.hatenablog.com
今回はclassy visionで学習したあとのモデルを読み込む処理を作成します
目次
スポンサーリンク
この記事でわかること
classy visionで学習したモデルを読み込む方法
1.実行環境
Jetson Xavier NX
ubuntu18.04
docker 19.03.6
Python 3.6.9
2.コード
from classy_vision.models import ClassyModel import torch parser = argparse.ArgumentParser(description='classy_vision') parser.add_argument('--pretrained_path') def load_ClassyModel(path:str) state = torch.load(path) model = ClassyModel.from_checkpoint(state) return model def main(): args = parser.parse_args() model = load_ClassyModel(args.pretrained_path) if __name__ == '__main__': main()
学習時に出力されるcheckpointファイルをtorch.loadで読み込んでから、
読み込んだ結果をClassyModel.from_checkpointに与えてると学習済みモデルを取得できます
スポンサーリンク
deeplearning関連記事
technoxs-stacker.hatenablog.com technoxs-stacker.hatenablog.com technoxs-stacker.hatenablog.com