[NVIDIA TAO] 画像分類モデル作成

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

進捗

前回やったこと

NVIDIA TAO（Train, Adapt, and Optimize）の画像分類モデル（Image Classification）を使って訓練。
訓練したモデルを評価。

今回やること

訓練したモデルで推論を実行。

現在のディレクトリ構造

/workspace/work
├ 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
│ │ ├ 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
│ └ tlt_pretrained_classification_vresnet18
│    └ resnet_18.hdf5
├ output
│ └ classmap.json ← 推論に必要なclassのindexとlabelをマップするファイル
│ └ training.csv
│ └ weights
│    ├ resnet_001.tlt
│    ├ resnet_002.tlt
│      …
│    └ resnet_010.tlt ← 推論で使用する訓練モデル
└ specs
   └ spec.cfg

classmap.jsonの確認

推論を実行するために必要なファイルで、訓練に成功すると $OUTPUT_DIR 直下に出力されます。

中を見てみると、classのindexとlabelが紐づけられたjson形式であることがわかります。

$ cat $OUTPUT_DIR/classmap.json
{
    "all": 0,
    "bassist": 1,
    "drummer": 2,
    "guitarist": 3,
    "vocalist": 4
}

推論を実行

下記のコマンドで推論を実行します。

※このコマンドはv2.0の場合のコマンドです。最新verのコマンドは公式ドキュメントで確認してください。

$ tlt-infer classification -m $OUTPUT_DIR/weights/resnet_010.tlt -i $DATA_DIR/test/drummer/201.jpg -cm $OUTPUT_DIR/classmap.json -k $KEY

オプションの説明

下記のオプション以外にもオプションがあるので興味があればこちらで確認できます。

オプション	説明
-m	訓練で出力されたモデル（.tltファイル）を指定します。
-i	推論対象の画像ファイルを指定します。
-d	推論対象の画像が格納されたディレクトリを指定します。
-cm	classのindexとlabelが紐づけられたjsonファイルを指定します。このファイルは訓練を実行すると自動出力されます。

実行結果

ドラマー（私）の写真を推論してみると、ちゃんと「drummer」に分類されました！

終わりに

これでNVIDIA TAO（Train, Adapt, and Optimize）の画像分類モデル（Image Classification）を使って、モデルの訓練～推論を実行して、画像分類を実行できることが確認できました！

これで大量にあるライブの写真を自動分類することができます！

本来NVIDIA TAO（Train, Adapt, and Optimize）はエッジデバイスやT4サーバにデプロイして使うのですが、私の目的ではローカルで推論して写真をフォルダ分けするshellを書けば事足りますからこれで良しとしましょうかね。

なんちゃってSEの足跡