Step-by-StepでAIを開発することのできるNVIDIA TAO Toolkitを利用して、AIモデルの入手から、学習、適応、最適化、デプロイまでの手順を理解する

MediaTekのSoCについて調べていたところ「NVIDIA TAO Toolkit」という開発ツールに出会いました。NVIDIA TAO Toolkit は 転移学習(Transfer Learning) によりAI開発の T(Train)、A(Adapt)、O(Optimize) を加速し、少ない手順で実運用に耐えるAIモデルを生成できる開発プラットフォームです。少ない手順で開発できるということは「NVIDIA TAO Toolkitによる開発手順さえ理解してしまえば、AI開発の基礎を習得できるのでは...」 と思い立ち、手元のワークステーションにNVIDIA TAO Toolkitをインストールし、AIモデルを開発してみました。

NVIDIA TAOによるAIモデルの開発が簡単な理由

NVIDIA TAO ToolkitによりAIモデルの開発手順を簡単に習得できる最大の理由は チュートリアルのソースコードに「Jupyter Notebook」を活用していることにあります。 Jupyter Notebookは、Pythonスクリプトに追加の説明文を追記できる形式の開発環境であり、 開発の説明文を読みながら、Pythonスクリプトによる開発を推進することができます。 つまり、NVIDIA TAO Toolkitの提供するJupyter Notebook用のファイル 「*.ipynb」 を上から順に辿り、それぞれのコードブロックの意味を理解し、実行していけば、Step-by-StepでAIモデルの開発方法を習得できてしまうのです。

NVIDIA TAO Toolkitのインストール手順

まず、NVIDIA TAO Toolkitをインストールしましょう。インストール手順は下記の記事をご参照ください。 NVIDIA TAO Toolkitは複数の開発フレームワーク上で動作しますので、依存するフレームワークをすべてインストールした後に、NVIDIA TAO Toolkitをインストールします。やや煩雑な手順となりますが、一度セットアップしてしまえば、利用する方法はとても単純なので、インストールに一度だけお付き合いください。

CUDA ... GPUへアクセスするためにインストール
cuDNN ... Neural NetworkからCUDAへアクセスするためにインストール
Docker ... NVIDIA TAOを実行するための仮想環境として追加
NVIDIA Container Toolkit ... Docker上からGPUを利用するために必要
NGC CLI ... NGC (NVIDIA GPU Cloud) を連携するためのCLIツール
Anaconda3 ... NVIDIA TAOと連携するための仮想環境も作成する

インストールしたNVIDIA TAOでYOLOv4を開発する

それでは、早速インストールしたNVIDIA TAO Toolkitを使って、 物体検出を行うAIモデル「YOLOv4-tiny」 を開発してみましょう。本記事の手順を辿ることにより、以下のように 自動車と人物を検出できるAIモデル を開発することが可能です。

実行環境

今回は、以下の環境にて、NVIDIA TAO Toolkitを利用しました。GPUのメモリサイズがSystem Requirenmentsに足りていませんが、 GPUのモデルサイズとして16GBが求められているのは、1回の推論で処理するデータの個数（batch_size）が16GB向けにチューニングされているため ですので、Memory Allocation Error等が発生してしまい、動作しない場合はbatch_sizeを調整すれば問題ないでしょう。

Core i5-14500 (14cores, 20threads)
RAM 64GB
NVIDIA GeForce RTX3060Ti 8GB x2
- 注意）System Requirenments(=16GB)に満たないメモリサイズ

NVIDIA TAOのSystem Requirementsは以下の通りです。

System Requirements
- 16 GB system RAM
- 16 GB of GPU RAM
- 8 core CPU
- 1 NVIDIA GPU
- 100 GB of SSD space

開発手順

NVIDIA TAO ToolkitによるAIモデルの開発は、
下記に挙げる11個の手順から構成されています。

Install the TAO launcher ... 環境設定
Prepare dataset and pre-trained model ... 学習済みAIモデルの取得
Provide training specification ... ハイパーパラメータの設定
Run TAO training ... 最初の転移学習
Evaluate trained models ... 学習結果の評価
Prune trained models ... 最適化（プルーニング）
Retrain pruned models ... プルーニングしたモデルでの学習
Evaluate retrained model ... 学習結果の評価
Visualize inferences ... 実画像に対する推論
Model Export ... AIモデルの書き出し
Verify the deployed model ... 書き出したモデルの評価

NVIDIA TAO Toolkitによる転移学習とは！

NVIDIA TAO ToolkitでYOLOv4のAIモデルを開発する

それでは、NVIDIA TAO Toolkitを使ってYOLOv4のAIモデルを開発してみましょう。

作業用ディレクトリを作成する

まず、NVIDIA TAO Toolkit用の作業用ディレクトリを作成します。

mkdir -p ~/sandbox_tao_tutorial/

Pythonの仮想環境を有効化する

次に、NVIDIA TAO Toolkitを実行するために必要なPythonの仮想環境「launcher」をactivateします。その後、 tao_tutorials ディレクトリにて jupyter notebook を実行し、ブラウザからNVIDIA TAOへアクセスできるようにします。 launcherのセットアップ方法については下記の記事をご参照ください。

(base) $ conda activate launcher
(launcher) $ cd ~
(launcher) $ cd tao_tutorials/
(launcher) $ jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root

# [I 2025-02-08 12:28:29.403 ServerApp] http://shino-Ubu22-devel:8888/tree?token=d0943ab5b36f7e39133dfef8c902bd8afe07c560c883a0c4
# [I 2025-02-08 12:28:29.403 ServerApp]     http://127.0.0.1:8888/tree?token=d0943ab5b36f7e39133dfef8c902bd8afe07c560c883a0c4
# [I 2025-02-08 12:28:29.403 ServerApp] Use Control-C to stop this server and shut down all kernels

NVIDIA TAO Toolkit内のYOLOv4-tiny開発スクリプトを開く

インストールしたNVIDIA TAO Toolkitに含まれる「YOLOv4-tiny」の開発スクリプトは 「notebooks」「tao_launcher_starter_kit」「yolov4-tiny」 に含まれています。

Jupyter Noptebook形式のYOLOv4-tinyの開発スクリプトを開くと、下記のように、説明文とソースコードの混在した開発画面が表示されます。この説明文を読みながら、 都度、作業ディレクトリなどの必要な情報を記入 し、Jupyter Notebookを実行していくことにより、YOLOv4-tinyのAIモデルを入手することができます。

作業用ディレクトリを作成し、指定する

続いて、 YOLOv4-tinyモデルの開発に利用するディレクトリを作成 します。そして、作成したディレクトリのパスをJupyter Notebookで開かれているYOLOv4-tinyの開発スクリプト中に記入します。 パスを指定する場合は、相対パスでなく、絶対パスを指定してください。

(base) $ mkdir -p ~/sandbox_tao_tutorial/yolov4_tiny
(base) $ cd ~/sandbox_tao_tutorial/yolov4_tiny/
(base) $ pwd
/home/shino/sandbox_tao_tutorial/yolov4_tiny

学習用データセットを入手し、展開する

NVIDIA TAO Toolkitは学習済みのAIモデルに対して、転移学習を行うことのできるツールキットです。ここでは転移学習に利用するデータセットをダウンロードします。

YOLOv4-tinyで利用する追加のデータセットは、 Notebookに記載されているリンクをクリックし、メールアドレスを入力することにより、メールアドレス宛にダウンロードURLが送られてきますので、そちらのURLをスクリプト中に記入することにより、入手できます。

データセットをダウンロードした後、TAOのコマンド tao model yolo_v4_tiny dataset_convert を利用して学習用データセットと検証用データセットを定義します。

学習済みモデルを入手する

次のステップを実行することにより、学習済みのAIモデルを入手できます。スクリプトの実行後 「*.hdf5」のTensorFlow 形式のAIモデルがダウンロードされた ことを確認できます。

転移学習を実行する

学習用データセットと、学習済みAIモデルを利用し、転移学習（transfer learning） により、目指すべきAIモデルを手早く入手します。 転移学習は tao model yolo_v4_tiny train コマンドにより実行し、ハイパーパラメータは *.kitti.txt にて指定 します。また、今回の開発環境には GPUが2基搭載されていますので --gpus 2 --gpu_index 0 1 を指定し両方のGPUを使う旨 を設定します。

学習が始まると nvidia-smi の出力にて、以下のように2基のGPUのメモリが消費されていることを確認できます（ 並列に学習が進んでいることを確認することができます ）。学習は80epochs行われます。

学習が完了したAIモデルを評価する

AIの転移学習が完了したならば、次に、AIモデルの評価を行います。 評価には tao model yolo_v4_tiny evaluate コマンドを利用します。 この評価により、どのくらいの精度のモデルが生成できたかを確認することが可能です。もし、十分に学習できていないようであれば 「学習のepochs」 や 「データセットの量や質」 を見直す必要があるでしょう。

影響の少ない重み値(weight)に対してプルーニング（枝刈り）を実施し、プルーニング後のモデルに再度学習を適用する

YOLOv4-tinyは非常に深いニューラルネットワークの構造を持ったAIモデルです。そのため、前手順により学習をした結果、接続しなくてもよい（重み値：weightを"0"として扱ってもよい）ニューロン間の結合が存在し、こうしたネットワークを切断することにより、AIモデルのレイテンシを改善することができます。このように、ニューラルネットワークの枝を刈る処理のことを 「プルーニング」 と呼びます。

プルーニングによる最適化では、まず、前述した通常の転移学習した結果に基づいて、不要と思われる枝を選定し、これを切断。新しくなったニューラルネットワークがよりよいパラメータを保持できるよう、再度学習を行います。

以下のコマンドにより プルーニングしたAIモデルに対して学習を再度実施します。

最後に、得られたAIモデルを評価しましょう。

学習した結果として得られたAIモデルで推論する

ここまでで得られたAIモデルを利用して、実際の画像データに対して推論を実施し、推論ができることを確認してみましょう。 推論には tao model yolo_v4_tiny inference コマンドを利用します。ここでは9枚の画像に対して推論を実施しています。 推論後、実際に画像を開くと、期待するオブジェクトがバウンディングボックスで囲まれている様子を見ることができます。

生成されたモデルから、FP32/FP16/INT8それぞれの重み値に量子化されたモデルを生成し、ONNX形式とTensorRT形式でエクスポートする

ここまでの手順で 構築したAIモデルを、実際のエッジAIでも実行できるようにデプロイ します。デプロイ工程では、量子化のデータ幅を決めるパラメータ fp32 fp16 int8 を tao model yolo_v4_tiny export や tao model yolo_v4_tiny gen_trt_engine に指定し、 エッジでも処理しやすいデータ幅(FP32/FP16/INT8)のAIモデルをエクスポートします。

エクスポートの結果、 FP32/FP16/INT8のAIモデルを得ることに成功 しました。

### 作業用ディレクトリへ移動する
(base) $ cd ~/sandbox_tao_tutorial/yolov4_tiny/yolo_v4_tiny
(base) $ pwd
# /home/shino/sandbox_tao_tutorial/yolov4_tiny/yolo_v4_tiny

### モデルがエクスポートされたディレクトリの内容を確認する
(base) $ ls -al ./export
# drwxrwxr-x  2 shino shino      4096  2月 10 15:32 .
# drwxrwxr-x 10 shino shino      4096  2月 10 15:12 ..
# -rw-r--r--  1 root  root       4746  2月 10 15:21 cal.bin
# -rw-r--r--  1 root  root  311869608  2月 10 15:20 cal.tensorfile
# -rw-r--r--  1 root  root         23  2月 10 15:12 labels.txt
# -rw-r--r--  1 root  root        193  2月 10 15:12 nvinfer_config.txt
# -rw-r--r--  1 root  root        795  2月 10 15:32 status.json
# -rw-r--r--  1 root  root   31962692  2月 10 15:14 trt.engine
# -rw-r--r--  1 root  root   13713724  2月 10 15:20 trt.engine.fp16
# -rw-r--r--  1 root  root    8707356  2月 10 15:32 trt.engine.int8
# -rw-r--r--  1 root  root   23706382  2月 10 15:12 yolov4_cspdarknet_tiny_epoch_080.onnx

ONNX形式のAIモデルの内容をNetronで確認する

生成されたONNX形式のネットワークは、Netronで開くことができます。YOLOv4-tinyは単純な深いネットワークですが、実際に開いてみると、プルーニングによりかなり複雑なネットワーク構造となっていることを確認できます。複雑なネットワーク構成となっているのは、枝刈りの影響でネットワークの分流が発生しているためです。

エクスポートしたTensorRT形式のAIモデルで推論を評価する

最後に、エクスポートしたTensorRT形式のAIモデルを利用し、推論できることを確認してみましょう。 trt.engine を引数に指定した tao deploy yolo_v4_tiny inference を実行 することにより、デプロイ後のAIモデルを使って推論することができます。

以下のように、 自動車と人物を検出できるAIモデルを得ることができました！

以上のように、NVIDIA TAO Toolkitを使うことにより、AIモデルを開発するフローをひとつずつ理解し、体系立ててAIモデルの開発手順を学ぶことができます。 AIモデルの取得から、学習、最適化、プルーニング、量子化までをひとつのスクリプトで扱っている例はネットを探してもほとんど見つけられない ため、是非一度NVIDIA TAO Toolkitをお試しください！

ここまで読んでNVIDIA TAO Toolkitをセットアップしてみたくなった方は以下の記事をご参照ください。ある程度のメモリとGPUを搭載したLinuxマシンがあれば、ローカルPCを使って無料で試すことができます。

NVIDIA TAO Toolkitでclassification(tf2)を開発する

次に yolo_v4_tiny の対比として、
NVIDIA TAO Toolkitで classfication(tf2) を開発する手順を掲載します。

開発の前準備を行う

まず、NVIDIA TAO Toolkit全体向けの作業用ディレクトリを作成します。

mkdir -p ~/sandbox_tao_tutorial/

次に、NVIDIA TAO Toolkit起動用のPythonの仮想環境をActivateします。

(base) $ conda activate launcher
(launcher) $ cd ~
(launcher) $ cd tao_tutorials/
(launcher) $ jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root

開発に利用するスクリプトを開く

今回開発するclassification(tf2)の開発用スクリプトは notebooks tao_launcher_starter_kit classification_tf2 tao_voc classification.ipynb に格納されています。

NVIDIA TAO x classfication(tf2) の notebook

classification(tf2) のnotebookは下記のようになっています。

開発時に利用するコマンドや引数は、下記のサイトにまとめられています。

開発環境の設定を行う

開発用のディレクトリを作成し、その絶対パスを、スクリプトへ記入します。

(base) $ cd sandbox_tao_tutorial/
(base) $ ls
# tao_ssd  yolov4_tiny

(base) $ mkdir classification_tf2
(base) $ cd classification_tf2
(base) $ pwd
# /home/shino/sandbox_tao_tutorial/classification_tf2

LOCAL_PROJECT_DIR がローカルPC上の開発用ディレクトリへのパスです。 その他の /workspace から始まるパスは、開発用のDockerコンテナの内部のディレクトリ構成を決めるパスです。これらのディレクトリは実行時に自動的に生成されますので、特に設定する必要はありません。

mountスクリプトを開き、開発用のディレクトリが指定されていることを確認します。

NVIDIA TAO Toolkitのパッケージを、ビルド環境にインストールします。

転移学習に利用するデータセットをダウンロードする

転移学習に利用するデータセットのURLをコピーし、LOCAL_PROJECT_DIR 内に data ディレクトリを作成し、wget コマンドを使ってデータセットをダウンロードします。

# 作業ディレクトリ内にdataディレクトリを作成する
(base) $ pwd
# /home/shino/sandbox_tao_tutorial/classification_tf2
(base) $ mkdir data

# dataディレクトリ内に学習用データセットをダウンロードする
(base) $ cd data
(base) $ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

ダウンロード後、スクリプトを実行し、展開すると下記のディレクトリが作られます。

学習済みAIモデルを取得する

続いて、学習済みAIモデルを取得します。が、提供されている「classification(tf2)」の開発用スクリプトは、何故か ngc registry model list と ngc registry model download-version がコメントアウトされており、AIモデルが自動的にダウンロードされるようになっていません。ここではAIモデルをダウンロードして利用したいため、 これらのコメントアウトを解除し、コードを有効化した後に、コードブロックを実行し、モデルをダウンロードしてください。

学習用のハイパーパラメータを設定する

提供されている「classification(tf2)」の開発用スクリプトが保持する学習用のハイパーパラメータでは、16GB RAMを持つGPU向けに、 batch_size と batch_size_per_gpu が大きな値となっています。今回は、RTX3060Ti 8GBモデルでAIの学習を実行するため、設定されている batch_size と batch_size_per_gpu を変更し、 1回に処理するデータ量を絞り、GPUのメモリが溢れないように調整しました。

!mkdir -p $LOCAL_EXPERIMENT_DIR/output
!sed -i "s|RESULTSDIR|$USER_EXPERIMENT_DIR/output|g" $LOCAL_SPECS_DIR/spec.yaml
### 以下を追加
# GPU毎のバッチサイズと、合計バッチサイズを調整する
!sed -i "s|batch_size_per_gpu: 64|batch_size_per_gpu: 16|g" $LOCAL_SPECS_DIR/spec.yaml
!sed -i "s|batch_size: 256|batch_size: 64|g" $LOCAL_SPECS_DIR/spec.yaml
# 更新された設定スクリプトを表示する
!cat $LOCAL_SPECS_DIR/spec.yaml

ハイパーパラメータが変更されていることを確認してください。

TAOによるAIモデルの転移学習を実行する

また、今回のAI学習環境は、RTX3060Tiを2基搭載していますので num_gpus=2 を指定し、学習に複数のGPUが利用されるよう設定しました。指定した後にスクリプトを実行し、学習を開始します。

学習を開始した後、Jupyter Notebookを実行しているプロンプトとは別のプロンプトを開き nvidia-smi を実行し、GPUのメモリが十分に利用されており、GPUのメモリが溢れそうになっていないかを確認しておきましょう。

80epochsの学習が完了すると、結果が表示されます。

学習結果を確認する

評価用のスクリプトを生成します。

tao model * evaluate コマンドにより評価を実行すると、精度が出力されます。

プルーニングによる最適化を行い、その後再度学習を実行する

続いて、プルーニング（枝刈り）を行います。

プルーニング後に実行する再学習で利用するハイパーパラメータも、上記の学習時と同様、RTX3060Ti 8GB向けのサイズに調整します。

!mkdir -p $LOCAL_EXPERIMENT_DIR/output_retrain
!sed -i "s|RESULTSDIR|$USER_EXPERIMENT_DIR/output_retrain|g" $LOCAL_SPECS_DIR/spec_retrain.yaml
!sed -i "s|PRUNEDMODEL|$USER_EXPERIMENT_DIR/output/prune/model_th=0.68_eq=union.tlt|g" $LOCAL_SPECS_DIR/spec_retrain.yaml
### 以下を追加
# GPU毎のバッチサイズと、合計バッチサイズを調整する
!sed -i "s|batch_size_per_gpu: 64|batch_size_per_gpu: 16|g" $LOCAL_SPECS_DIR/spec.yaml
!sed -i "s|batch_size: 256|batch_size: 64|g" $LOCAL_SPECS_DIR/spec.yaml
# 更新した設定スクリプトを表示する
!cat $LOCAL_SPECS_DIR/spec_retrain.yaml

tao model * train を再度実行し、プルーニング後の再学習を行います。

Quantization Aware Training（QAT）向けのコードブロック

TensorFlowは、Model Optimization Toolkitとして 「Quantization Aware Training（QAT）API」 を提供しています。 QATは、ニューラルネットワークの学習中に低精度のハードウェアをシミュレートし、全体的なネットワーク損失メトリックに量子化誤差を最小化する機能です。NVIDIA TAO ToolkitからQATを利用するにはQATに関するコードブロックを実行してください（ 今回の手順では、QATの実行を省略します ）。

モデルを評価する

学習によって得られたモデルを tao model * evaluate により評価します。

生成できたAIモデルをエクスポートする

最後に生成したAIモデルをエクスポートしましょう。 tao model * export コマンドにより、ONNX形式としてAIモデルをエクスポートすることができます。 エクスポート時に指定できるパラメータについては下記のサイトをご参照ください。

エクスポートされたONNX形式のAIモデルは、Netronにより可視化できます。

エクスポートしたAIモデルにより推論できることを確認する

最後に、エクスポートしたAIモデルを指定した tao deploy * inference コマンドにより推論できることを確認しましょう。

以上が classfication(tf2) の開発の流れです。
yolo_v4_tiny とほぼ同じ手順で開発できることがわかります。

このようにNVIDIA TAO Toolkitを使うことで、様々なネットワークを体系立てて、転移学習、プルーニング、量子化といった流れをスムーズに実行することができます。

是非ご活用ください！