[NVIDIA TAO] 画像分類モデル作成 - 訓練編 [4/5]

NVIDIA TAO（Train, Adapt, and Optimize （旧名 NVIDIA TLT））で画像分類モデル（Image Classification）を作成～推論に挑戦する記事の４回目です。

進捗

1. 環境構築
2. 事前訓練モデルをダウンロード
3. 訓練データの作成
4. モデルの訓練　←　現在地
5. 推論

前回やったこと

• 訓練データを用意する。
• 設定ファイルを作成する。

今回やること

• 用意した訓練データ＆設定ファイルを使ってモデルを訓練する。

 

現在のディレクトリ構造

/workspace/work ← $WORK_DIR
├ data ← $DATA_DIR
│ ├ train ← 訓練で使用する画像
│ │ ├ drummer ← 任意の分類名
│ │ │ ├ 001.jpg
│ │ │ └ 002.jpg
│ │ ├ guitarist
│ │ │ ├ 003.jpg
│ │ │ └ 004.jpg
│ │ ├ bassist
│ │ │ ├ 005.jpg
│ │ │ └ 006.jpg
│ │ ├ vocalist
│ │ │ ├ 007.jpg
│ │ │ └ 008.jpg
│ │ └ all
│ │    ├ 009.jpg
│ │    └ 010.jpg
│ ├ val ← 検証（Validate）で使用する画像
│ │ ├ drummer
│ │ │ ├ 101.jpg
│ │ │ └ 102.jpg
│ │ ├ guitarist
│ │ │ ├ 103.jpg
│ │ │ └ 104.jpg
│ │ ├ vocalist
│ │ │ ├ 107.jpg
│ │ │ └ 108.jpg
│ │ └ all
│ │    ├ 109.jpg
│ │    └ 110.jpg
│ └ test ← テストで使用する画像
│    ├ drummer
│    │ ├ 201.jpg
│    │ └ 202.jpg
│    ├ guitarist
│    │ ├ 203.jpg
│    │ └ 204.jpg
│    ├ bassist
│    │ ├ 205.jpg
│    │ └ 206.jpg
│    ├ vocalist
│    │ ├ 207.jpg
│    │ └ 208.jpg
│    └ all
│       ├ 209.jpg
│       └ 210.jpg
├ model ← $MODEL_DIR
│ └ tlt_pretrained_classification_vresnet18
│    └ resnet_18.hdf5 ← 事前訓練モデル
├ output ← $OUTPUT_DIR
└ specs ← $SPECS_DIR
   └ spec.cfg ← 設定ファイル

モデルの訓練

このコマンドが頻繁に変わるの、ホントやめてほしい...
tao classification だったり classification train だったり tlt classification train だったり↓だったり、少しバージョンが変わればコマンドが変わってきます。

$ tlt-train classification -e $SPECS_DIR/spec.cfg -r $OUTPUT_DIR -k $KEY

訓練にかかる時間は

• GPU … RTX2070Super
• インプット画像 … 625枚
• エポック数 … 10
• バッチサイズ … 40

で12分程度でした。

モデルの出力結果

訓練が終わると、$OUTPUT_DIR/weights/ にモデル（.tltファイル）が出力されます。

.tltファイルはエポック数の数だけ出力されます。

$ ll $OUTPUT_DIR/weights/

 

 

 

 

 

モデルの評価

訓練したモデル（.tltファイル）の評価を行います。

評価の設定は前回の記事のeval_configで定義しているので、下記のコマンドを実行するだけです。

$ tlt-evaluate classification -e $SPECS_DIR/spec.cfg -k $KEY

終わりに

これでNVIDIA TAO（Train, Adapt, and Optimize）で画像分類モデル（Image Classification）の訓練モデルを作成できました。

次回はいよいよ推論を実行してみます。