これまで画像分類モデルの学習をさせる際にpytochのフレームワークのclassyvisionを使っていました
フレームワークは使いやすいのですが、自分で拡張することが難しそうです・・・
フレームワークでなくテンプレートであれば色々拡張できそうでしたので、今回よさそうなテンプレートを探して、試してみました
この記事は備忘録になります

この記事でわかること

pytorchテンプレートを使った学習環境を作る方法

1.環境構築

1.1 実行環境

Jetson Xavier NX
ubuntu18.04
docker
python3.6

1.2 テンプレートの取得

今回使用したテンプレートは以下になります
github.com

以下コマンドでgithubからcloneします

cd workspace
git clone https://github.com/victoresque/pytorch-template.git

1.3 dockerfile&imageの作成

今回使用したdockerfileは以下になります

FROM nvcr.io/nvidia/l4t-pytorch:r32.4.4-pth1.6-py3

RUN pip3 install --upgrade pip
RUN pip3 install --ignore-installed PyYAML
RUN pip3 install tensorboard
RUN pip3 install pandas

こちらをbuildします

sudo docker build . -t pytorch_templete

1.4 docker containerの起動

docker containerを起動して、コンテナ内に入ります

sudo docker run -it --rm --runtime nvidia -v /path/to/your/workspace/dir/:/workspace --workdir /workspace --network host pytorch_templete

2.サンプル実行

今回はmnistとresnetによる画像分類モデルの学習を実行してみます
config.jsonにはデフォルトで上記内容の学習になるように記述されているので、このまま使います

python3 train.py -c config.json

実行後の評価結果は以下のような感じです

    epoch          : 57
    loss           : 0.07924625007832926
    accuracy       : 0.9757558606973595
    top_k_acc      : 0.9974081753554502
    val_loss       : 0.03324085925061731
    val_accuracy   : 0.9901690729483283
    val_top_k_acc  : 0.999501329787234

テンプレートを使った学習ができるようになりました

感想

シンプルな構造のテンプレートなので、拡張しやすそうです

3.参考

github.com

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

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

【pytorch】pytorchテンプレートを使った学習環境を作る方法

目次