Vitis™ AIを用いたPyTorchモデルの実装フローをYOLOX用いてご紹介！<準備編>

2024/04/19

皆さまこんにちは！TEPPE-AIです！
大変ご無沙汰しております。

今回ご紹介するものは
「Vitis™ AIを用いたPyTorchモデルの実装フローをYOLOX用いてご紹介！<準備編>」です。

Vitis™ AIでは、現時点でTensorFlow, PyTorch（1.3よりサポート開始）, ONNX（3.5より新サポート）のフレームワークに対応しています。
その中で今では設計界隈ではおなじみのPyTorchにフォーカスして、実装フローを解説します！

今回は、準備編と題し、YOLOXという物体検出を行うAIモデルをKria™ KV260上（以降、KV260）に実装する手順を例に紹介していきますので、ぜひ最後までお付き合いいただけると嬉しいです。

それでは参りましょう！

Vitis™ AIについて

Vitis™ AIは、学習済みAIモデルをAMD社製SoCデバイス上で高速演算させる変換環境になります。

図1. Vitis™ AIについて①

Vitis™ AIは主に3つのツールから構成されており、最新バージョン（3.5）以降はすべて無償で利用可能となっています。

AI Quantizer（量子化）：学習済みモデルをFP32からINT8へ変換
AI Compiler（コンパイル）：量子化済みモデルを最適化
AI Optimizer（プルーニング）：過剰な重みを削除し、演算コストを削減　*3.5以降無償化

図2. Vitis™ AIについて②

Vitis™ AIによる変換は主に2工程実施され、「AI Quantizerによる量子化→AI Compilerによるコンパイル」です。
AI Quantizerにより学習済みモデルが量子化（FP32→INT8）され、AI Compilerにより搭載しているDPU向けに最適化されます。

AI Optimizerは、一般的に言われるプルーニングを行う機能で、認識精度に影響のないノードを削除し、演算コストを削減し、処理速度の向上が見込めます。
使用は任意で、処理速度が向上しない際の検討候補として用いることをおすすめします。

なおVitis™ AIは2019/11リリース以降、定期的なアップデートを行っており、2024年3月時点での最新バージョンは3.5になります。

図3. Vitis™ AIアップデート推移

今回は、PyTorchベースの学習済みモデルの実装フローにフォーカスしていきます。
流れとしては、AI Quantizer（Pythonモジュール）にて量子化、AI Compiler（Linuxコマンド）を用いて変換を行います。

図4. Vitis™ AI PyTorchベース実装フロー概要

作業概要

Vitis AIによるAIモデル変換含め、KV260への実装までを行います。概要は以下の通りです。

図5. 作業概要

用いる環境は、以下の通りです。

図6. 使用デバイス環境

Vitis AIは最新の3.5ではなく、3.0を用います。
3.5はマイナーアップデートとなっており、ターゲットデバイスとしてVersal AI Edgeシリーズとなります。
用いるKV260（Zynq™ UltraScale+ ™ MPSoC）は対象外となっています。

※: Vitis AI 3.5のアップデートサマリは、ブログ記事として公開しておりますので、こちらをご確認ください！

ただし、Vitis AI 3.5環境はサポートしているので、

Vitis AI 3.0環境を用いること
Vitis AI 3.5のHWデザインをSoftwareツール（Vivado, Petalinux, Vitis）を使って構築

上記いずれかで対応するようAMD公式からのアナウンスが記載されていました。
詳細はこちらに記載されています。あわせてご確認ください。

今回はPyTorchモデルの実装フロー紹介をメインとしたく、リファレンスデザインが用意されているVitis AI 3.0環境のを用いることにしました。
KV260のVitis AI 3.5環境構築方法についても併行して検証中ですので、完了次第、記事にしたいと思います！

図7. Vitis AI 3.5アップデートサポートデバイスについて

出典：AMD Supported Evaluation Targets（参照 2024/4/4）

なお今回、学習工程は含みません。
こちらのGithubからダウンロードできる学習済みのYOLOXを活用します。
サンプルを活用し、実装までの流れを紹介することをメインとさせていただきます。
是非ご参照いただき、皆さまの環境に合わせてサンプルを改変いただければと思います。

使用機材

左右にスクロールしてご覧ください

hostPC	OS	Ubuntu 22.04 LTS	Linux環境が必須となりますのでご注意ください。 GPUは必須ではありません。
hostPC		Ubuntu 22.04 LTS	作業項目	②各種ダウンロード ③学習済みAIモデル量子化/コンパイル
Kria™ KV260 ビジョン AI スターターキット		Linux（yocto）	今回はリファレンスデザインを用いて検証を行います。
Kria™ KV260 ビジョン AI スターターキット		Linux（yocto）	作業項目	④推論アプリ構築、⑤実機確認
KV260操作用PC		Windows10	USB-UARTでKV260を操作しますので、操作用のソフト、ドライバのインストールをお願いします。
KV260操作用PC		Windows10	作業項目	KV260操作
microSDカード	Micron製 industrial micro SD card 32GB		リファレンスデザインはSD IMGのため、起動用として用意します。 PALTEKで取り扱っておりますので、気になる方はこちらからクリックお願いします！

図8. 接続イメージ

デバイス環境構築

hostPC

Vitis AIダウンロード

以下コマンドでダウンロードします。

左右にスクロールしてご覧ください

$ cd ~/Vitis/vitis_r3.0
$ git clone -b 3.0 https://github.com/Xilinx/Vitis-AI

Vitis AI dockerイメージインストール

変換環境として、dockerを用います。
CPUとGPUで用意がありますが、本作業ではどちらでも大丈夫です。

＜CPUの場合＞

Docker hubにビルド済みのDockerイメージがあるので、そちらを活用します。
上記のリンクにアクセスしたのち、今回はPyTorchベースなので
以下の「xilinx/vitis-ai-pytorch-cpu」をクリックし、インストールコマンドを確認します。

図9. Docker hub AMD(xilinx)提供Docker imageページ

出典：dockerhub AECG（参照 2024/4/4）

クリックし、「Tags」を選択すると、Vitis AIのバージョンごとにコマンドが記述されています。
“TAG：latest”でも問題ありませんが、バージョンごとに保存することを推奨しています。
今回は、以下ハイライトにある”TAG:ubuntu2004-3.0.0.106”のコマンドを用います。

※: 開発マシン（hostPC）のOSはUbuntu22.04ですが、Docker上ではUbuntu20.04になります。
作業にあたっての支障はないのでご安心ください！

図10. Docker hub PyTorch版Docker image一覧

出典：dockerhub　xilinx/vitis-ai-pytorch-cpu（参照 2024/4/4）

左右にスクロールしてご覧ください

$ docker pull xilinx/vitis-ai-pytorch-cpu:ubuntu2004-3.0.0.106

インストール完了したのち、以下コマンドでDockerイメージがあることを確認します。

左右にスクロールしてご覧ください

$ docker image ls

※: 権限がないとエラーが出た際はsudoをつけるか、dockerグループ作成し追加する作業を実施してください。【1】

図11. docker image ls出力画面

REPOSITORY：xilinx/vitis-ai-pytorch-cpu　TAG：ubuntu2004-3.0.0.106
が表示されていればインストール完了です！

＜GPUの場合＞

GPU版はユーザー側でビルドを行います。
Vitis AI Githubリポジトリにてビルド用のスクリプトがあるので、そちらを活用します。
Vitis-AI/dockerディレクトリ内のdocker_build.shが該当のスクリプトになります。
各種オプションがあるので、実行する際はこちらをあわせてご参照ください。

左右にスクロールしてご覧ください

$ cd ~/Vitis/vitis_r3.0/Vitis-AI/docker
$ ./docker_build.sh -t GPU -f pytorch

実行完了後、CPU版と同様、ビルド完了しているか確認します。
REPOSITORY：xilinx/vitis-ai-pytorch-gpu　TAG：3.0.0.001
が表示されていれば、ビルド完了です！

図12. docker image出力画面

評価ボード

リファレンスデザイン SD IMG書き込み

GitHub上にあるリファレンスデザインをダウンロードします。
リファレンスデザインには、OSやDPU含め、動作に必要なライブラリ等がすべて含まれています。
以下ハイライト箇所をクリックすると、ダウンロードされます。

図13. Vitis AIリファレンスデザイン格納先

出典：AMD Quick Start Guide for Zynq™ UltraScale+™（参照 2024/4/4）

インストール後、SD書き込みツール等を用いてmicroSDカードに書き込んで準備完了です！

各種データダウンロード

学習済みAIモデル

こちらのYOLOXを活用します。

図14. YOLOX公式GitHub

出典：Megvii-BaseDetection（参照 2024/4/4）

README.mdをスクロールすると、重みファイルが格納されていますので、
以下にあるYOLOX-nanoの重みファイルをダウンロードします。

図15. YOLOX GitHub 重みダウンロードページ

出典：Megvii-BaseDetection（参照 2024/4/4）

クリックするとダウンロードできますが、今回もwgetコマンドを用いてダウンロードします。

左右にスクロールしてご覧ください

$ cd ~/Vitis/vitis_r3.0/Vitis-AI/model_zoo/pt _yolox-nano_coco_416_416_1G_3.0/code/float
$ wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.pth -O yolox_nano.pth

キャリブレーションデータ

量子化時に必要なデータセットをダウンロードします。
量子化した際の精度を担保するためのキャリブレーションとして用います。
ケースバイケースになりますが、必要量として100~1000枚の画像で、ラベルは不要になります。

図16. AI Quantizerフロー概要図

出典：AMD Technical Information Portal Vitis AI ユーザーガイド(UG1414)（参照 2024/4/4）

今回ダウンロードしたAIモデルはCOCO2017データセットを学習しているため、
COCO2017の評価用データセット（val2017）をダウンロードします。

図17. COCOデータセット公式ページ

出典：COCO Common Objects　in Comtext（参照 2024/4/4）

val2017.zipをダウンロードし、data/COCOディレクトリに解凍していきます。

左右にスクロールしてご覧ください

$ cd ~/Vitis/vitis_r3.0/Vitis-AI/model_zoo/pt _yolox-nano_coco_416_416_1G_3.0/code/data
$ wget http://images.cocodataset.org/zips/val2017.zip -O val2017.zip
$ unzip val2017.zip -d COCO/

以下のようになればOKです！

図18. COCO2017データセットディレクトリ

量子化スクリプト

Vitis AIにてAIモデルを変換する際、量子化・コンパイルを行うスクリプトを作成する必要があります。
Vitis AI Model Zoo内に量子化スクリプトのサンプルがあるので、こちらを活用します。
こちらからVitis AI Model Zooのページにアクセスし、
「pt_yolox-nano_coco_416_416_1G_3.0」をクリックします。

図19. Vitis AI Model Zoo(3.0)ページ

出典：Xilinx/Vitis-AI（参照 2024/4/4）

クリックすると、以下のように表示されるので、
「https://www.xilinx.com/bin/public/openDownload?filename=pt_yolox-nano_coco_416_416_1G_3.0.zip」にアクセスしダウンロードするか、wgetコマンド等のコマンドでダウンロードを行います。

図20. YOLOXサンプルページ

出典：Xilinx/Vitis-AI（参照 2024/4/4）

今回はwgetを用いてダウンロードを行います。
ダウンロードし、unzipコマンドで解凍できれば作業完了です。

左右にスクロールしてご覧ください

$ cd ~/Vitis/vitis_r3.0/Vitis-AI/model_zoo
$ wget https://www.xilinx.com/bin/public/openDownload?filename=pt_yolox-nano_coco_416_416_1G_3.0.zip -O pt_yolox-nano_coco_416_416_1G_3.0.zip
$ unzip pt_yolox-nano_coco_416_416_1G_3.0.zip

解凍後、以下のようにディレクトリが確認できればOKです。