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

はじめに

NVIDIA TAO（Train, Adapt, and Optimize（旧名 NVIDIA TLT））とはNVIDIAが提供する様々な事前訓練済みのモデルを使用して転移学習を行うことで、ユーザのニーズに合わせたAIモデルを構築する為のツールです。

事前訓練済みモデルには

物体検出（Object Detection）
画像分類（Image Classification）
意図解釈（NLU）
視線推定
テキスト認識

を含めて他大多数の種類が提供されており、それぞれのモデルをそのまま使うことも出来ますし、ユースケースに合わせて自由に訓練を行い、モデルをカスタマイズすることが出来ます。

今回はその中でも画像分類（Image Classification）のモデルを使って、訓練や推論に挑戦してみるのを全５つの記事に分けて書いていこうと思います。

目的

自分のライブで撮影された大量の写真をパート毎に分類したい。
AIを使って自動で分類したい。
ローカルで完結したい。（大量の画像が存在して、クラウドにアップロードなんてしてられない）
NVIDIAが好きなのでNVIDIA TAO（Train, Adapt, and Optimize）で画像分類モデル（Image Classification）を使ってみたい。

進捗

ハードウェア要件

ハードウェア	推奨
メモリ	32GB
CPU	8コア
GPU	NVIDIA A100, NVIDIA V100, NVIDIA RTX 30X0
ストレージ	100GB

ソフトウェア要件

ソフトウェア	バージョン
Ubuntu	18.04 LTS
docker-ce	>19.03.5
docker-API	1.40
nvidia-driver	>455

筆者の環境

NVIDIA TAOはUbuntu + Dockerでの利用が前提となっていますが、私はWindows11 + Docker（WSL2バックエンド）で動かします。

OS	Windows11 + WSL2（Ubuntu20.04 LTS バックエンド）
メモリ	32GB
CPU	Intel Core i7-10700（8コア）
GPU	NVIDIA RTX 2070 Super

WSL2で動かす場合の注意点１

WSL2でGPUを使用するためには、Windows11とWSL2の両方に設定が必要で結構手間なのですが、幸いにもNVIDIA様がQiitaに記事を載せてくれていて、基本コピペしていけばWSL2でGPUを使用出来るようになります。

ちなみに私はDocker DesktopとWSL2バックエンドを使う方法でやっていますが、参考記事内で書かれている通り、動きはするものの未サポートなのでオススメしません。

参考：https://qiita.com/ksasaki/items/ee864abd74f95fea1efa

WSL2で動かす場合の注意点２

2021/12/28現在、NVIDIA TAO（Train, Adapt, and Optimize）の最新版（v3.0）はWSL環境では動作しません。

というのも訓練でエラーが出て、それ以上先に進めなくなります。

WSL2で動かしたい場合、下記のissueに記載の通り、古いバージョン（v2.0）のDocker Imageをpullしてくる必要があります。

参考：https://forums.developer.nvidia.com/t/wsl2-tao-issues/195476

事前に必要なもの

画像分類（Image Classification）の事前訓練済みモデルをダウンロードするためにNGC APIキーが必要です。

まだ発行していない場合は NVIDIA NGC でアカウントを作成してNGC APIキーを発行しておきましょう。

環境構築

NVIDIA TAO（Train, Adapt, and Optimize）のDocker Imageを取得

※Windows11 + Docker(WSL2バックエンド)の場合はWSL2のターミナルではなくWindowsのPowerShellからコマンドを実行します。

WSL2の場合（古いバージョン）

$ docker pull nvcr.io/nvidia/tlt-streamanalytics:v2.0_py3

Ubuntuの場合（2021/12/28時点での最新）

$ docker pull nvcr.io/nvidia/tlt-streamanalytics:v3.0-py3

Dockerコンテナを起動

$ docker run --gpus all -it -v "{Dockerにマウントするホストのディレクトリ}":"{Docker内のマウント場所}" nvcr.io/nvidia/tlt-streamanalytics:v2.0_py3

例えば、下記のように実行するとWindowsの `D:/tao` とDocker内の `/workspace` が同じファイルを参照できるようになります。（WSL2のパスを指定してもマウントされなかったので、Windowsのパスを指定）

$ docker run --gpus all -it -v D:\tao:/workspace nvcr.io/nvidia/tlt-streamanalytics:v2.0_py3

NGCのコンフィグ設定

事前訓練済みモデルをダウンロードするための設定を行います。
下記のコマンドを実行したら表示されるメッセージに従って事前に作成したNGC APIキーを設定してください。

$ ngc config set

環境変数の設定

$ export KEY={NGC APIキー}
$ export WORK_DIR=/workspace/work
$ export MODEL_DIR=$WORK_DIR/model
$ export DATA_DIR=$WORK_DIR/data
$ export SPECS_DIR=$WORK_DIR/specs
$ export OUTPUT_DIR=$WORK_DIR/output

作業ディレクトリ作成

$ mkdir -p $WORK_DIR
$ mkdir -p $MODEL_DIR
$ mkdir -p $DATA_DIR
$ mkdir -p $SPECS_DIR
$ mkdir -p $OUTPUT_DIR

終わりに

これでNVIDIA TAO（Train, Adapt, and Optimize）を利用するための環境が整いました。

次回はこの環境に画像分類（Image Classification）の事前訓練済みモデルをダウンロードしていきます。

なんちゃってSEの足跡