NVIDIA TAO Toolkitを使ってみた

今回は、NVIDIA Train Adapt Optimize (TAO) Toolkitを実際に使用した体験やその魅力について紹介します。TAO Toolkitは事前学習済みチェックポイントや専用ツールキットを使用して、簡単にAI活用をはじめることができるツールキットです。これからAI活用をはじめたいと考えている方はぜひ参考にしてみてください。

TAO Toolkitとは

TAO Toolkitは、事前学習済みのAIモデルを取得し、ユーザー独自のデータでカスタマイズするための、シンプルかつ使いやすいAIツールキットです。

参照：https://developer.nvidia.com/ja-jp/tao-toolkit

まず、事前学習済みモデルにユーザーが準備する独自モデルやNVIDIAが提供するモデルを使用します(図①)。NVIDIAが提供するモデルにはタスクベースでカスタマイズされたものもあるため、ユーザーはAIを活用したい目的に合わせて直感的にモデル開発を始めることができます。次に、これらの事前学習済みモデルにユーザー独自のデータで転移学習を行い、モデルカスタマイズを行います(図②)。この学習は、TAO Toolkitに準備されているコマンドを使用し、簡単に実行することができます。AIの学習に必要な各種パラメータは設定ファイルに記述するため、ユーザーは複雑なコード入力などを行う必要がありません。最後に、カスタマイズされたモデルを推論環境向けに出力します(図③)。この際には、プルーニングや量子化を使用して、モデルサイズを調整することが可能です。以上のような手順で簡単にカスタムAIモデルを作成できるのがTAO Toolkitです。

TAO Toolkitの動作環境

まず初めに、簡単にTAO Toolkitの動作環境について紹介します。現在(2023年9月時点)のTAO Toolkitを使用するために必要なハードウェア要件は以下です。

最低限動かすために必要な要件は非常に少ないので、個人で所有するGPU搭載PCでも十分に動作させることが可能です。※TAO Toolkitに含まれる全ての機能を使用したい場合は、推奨要件以上の構成を準備しましょう。また、TAO ToolkitはAWSやAzureなどのクラウドGPUインスタンス上でも動作させることができます。本稿を読み、興味を持たれた方はぜひ自宅PCやクラウドで触ってみてください。今回、私がTAO Toolkitを動かした環境は、当社のTED AI Labに設置しているDGX A100です。

出典：NVIDIA

DGX A100は最新世代のA100 GPUを八枚搭載したアプライアンスサーバーで、AIモデル開発に必要な様々な計算を非常に高速に処理することができます。

次にTAO Toolkitの環境構築方法についてですが、TAO ToolkitはPythonベースのツールキットであるため、pipを使用して簡単にインストールすることが可能です。また、簡単にサンプルノートブックだけを動かしてみたいという方は、セットアップ用のスクリプトが準備されているので、そちらを使用すればすぐに動かすことができます。詳細な環境構築手順については、NVIDIAドキュメントをご確認ください。

サンプルノートブックで即座に体験

本稿では、TAO Toolkit内に準備されているサンプル用のノートブックを実際に動かしてみたいと思います。TAO Toolkitでは数多くのサンプルノートブックが配布されており、そのほとんどが下図のように実際のAI活用を想定した内容になっています。そのため、ユーザーはこのサンプルノートブックを動かすだけで、簡単かつ即座にAI活用をはじめることができます。

出典：NVIDIA

今回は、事前学習済みのResNet18をKITTIデータセットで転移学習し、自動車などの物体検出が可能なSSDモデルを作成するノートブックを、実際に動かしていこうと思います。サンプルノートブックは環境のセットアップから推論モデルの出力まで全てのコードが記載済みであるため、データセットやチェックポイントを配置するためのパスさえ設定すれば、ユーザーはコードを実行していくだけでTAO Toolkitを体験することが可能です。今回のノートブックで体験できるモデル作成手順を以下に紹介します。

環境のセットアップ
データセットと事前学習済みモデルの準備
学習時のハイパーパラメータ設定
事前学習済みモデルの転移学習
転移学習済みモデルの検証
転移学習済みモデルのプルーニング
プルーニングモデルの再学習
再学習済みモデルの検証
推論の可視化
モデルの出力

上記のように、AI開発の基本的な流れを体験できる上に、2~8までの手順のほとんどはTAO Toolkitのコマンド一つで実行することができます。実際に、私がこれらの手順を実行し、作成したモデルの推論結果が以下です。

図の通り、自動車を検出することができるモデルになっていると思います。また、このモデルは手順6にてプルーニングを行っているため、モデルサイズが12MB程度と小さく、推論環境として用いられるエッジデバイスでも十分に活用可能です。今回の体験では、前項にて紹介した通り、最新世代のA100 GPUで計算を加速したこともあり、環境セットアップからわずか2~3時間でこのような推論モデルを作成できました。「AI活用には多大なコストがかかる」と思っている方はまだまだ多いと思いますが、TAO ToolkitとAI向けGPGPUを活用すれば、即座にAI活用を体験できることがわかっていただけたと思います。

TAO Toolkitの魅力

TAO Toolkitを実際に動かしてみて、私がとくに便利だと感じたTAO Toolkitの魅力について三つほど紹介します。

- 豊富な事前学習済みモデル
  
  TAO Toolkitでは、数多くの事前学習済みモデルを提供しています。今回紹介した画像系のモデルのみでなく、自然言語系のモデルもサポートされているため、ユーザーはAIを活用したいタスクに合ったモデルをすぐに見つけることができます。
  
  ※出典：NVIDIA
- DALIとの連携
  
  TAO Toolkitは、同じくNVIDIAから提供されているデータ前処理のためのライブラリであるDALIと連携して使うことができます。DALIはデータ前処理をGPUで行うことでAI開発パイプライン全体を高速化することが可能です。また、TAO Toolkitでモデルカスタマイズを行う際に、カスタマイズのためのユーザー独自データの量が少ない場合、DALIを使用してデータ拡張を行うことができるため、よりモデルをカスタマイズしやすくなります。
  
  出典：NVIDIA
- 様々な推論環境に合わせたモデル出力
  
  TAO Toolkitでは、推論モデルの出力先となるエッジデバイスにあわせて、INT8やFP16などの様々な精度のデータタイプでモデル出力を行うことができます。また、最新バージョンのTAO ToolkitからはONNX形式でのモデル出力もサポートされるようになったため、GPUのみでなく、CPUやMCUなどの推論デバイスへもモデルデプロイが可能です。
  
  出典：NVIDIA

最後に

いかがだったでしょうか。本稿を読んで、TAO Toolkitを使ったAI活用をはじめてみたいと思っていただけたら嬉しいです。もし「TAO Toolkitを使ってみたいが環境が無い」、「設定方法や使い方がわからない」などの理由でお困りの場合には、ぜひ当社にご相談ください。当社では、TED AI Lab環境をお客様へお貸出しするTED AI Lab Engineering Service(TAILES)というサービスを提供させていただいております。TED AI Labには本稿で紹介したDGX A100などの最新世代のGPU環境を備えられており、存分にTAO Toolkitを試していただくことができます。また、単に機器をお貸出しするだけでなく、当社エンジニアが環境設定から使用方法まで手厚くサポートさせていただきます。

ブログをお読みになり、NVIDIA製品、TED AI Lab、TAILESにご興味がある方はぜひ当社までお問合せ下さい。