こすたろーんエンジニアの試行錯誤部屋

作成物の備忘録を書いていきますー

【classy vision】学習したモデルを読み込む方法

以前に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

参考

classyvision.ai