Vitis AI / ROCm の利用方法

2024.04.10

ACRi ルームでは AMD VCK5000 Versal 開発ボードを搭載する as005 サーバと、AMD Instinct MI210 アクセラレータを搭載する as006 サーバの無償利用環境を提供しています（詳細はこちら）。これらのサーバを利用することで、AI モデルの学習から Vitis AI によるモデルの量子化、コンパイル、実行までを行うことができます。

ROCm を使用した AI モデルの学習

ROCm は AMD GPU を使って AI などの計算を高速化するソフトウェアです。as006 サーバには ROCｍ がインストールされており、ROCｍ を利用した高速な AI モデルの学習ができます。

PyTorch のインストール方法

as006 に SSH またはリモートデスクトップでログインしターミナルを開きます。

作業ディレクトリ内に venv 環境を作成し、ROCm に対応した PyTorch をインストールするには次のようにします。

mkdir /scratch/$USER/my-pytorch-project
cd /scratch/$USER/my-pytorch-project
python -m venv .venv
source .venv/bin/activate
pip install --find-links=/tools/pip/pytorch torch==2.2.2+rocm5.7 torchvision==0.17.2+rocm5.7 torchaudio==2.2.2+rocm5.7

ROCm が有効であることを確認します。

python -c "import torch ; print(torch.cuda.is_available())"

あとは PyTorch を使用して AI モデルの学習を行うことができます。

ACRi ルームのサーバではインターネットへのアクセスができませんが、pip を使用した pypi.org からのパッケージのインストールはユーザーが自由にできるようになっています。例えば numpy をインストールするには次のようにします。

pip install numpy

PyTorch による画像分類チュートリアル

CIFAR10 データセットの画像分類を ResNet モデルで学習するチュートリアルを実行してみましょう。as006 にログインし、スクラッチ領域にチュートリアルのリポジトリをコピーします。

cd /scratch/$USER
cp -a /tools/repo/acri-room/vai-tutorial .
cd vai-tutorial

venv 環境を作成し PyTorch をインストールします。

python -m venv .venv
source .venv/bin/activate
pip install --find-links=/tools/pip/pytorch torch==2.2.2+rocm5.7 torchvision==0.17.2+rocm5.7 torchaudio==2.2.2+rocm5.7

他に必要となるパッケージをインストールします。

pip install -r requirements.txt

CIFAR10 データセットを data ディレクトリ以下に展開します。

mkdir data
tar xf /tools/data/cifar-10-python.tar.gz -C data

モデルを学習するには次のスクリプトを実行します。

./scripts/train.sh

デフォルト設定では200エポックの学習を行い、15分ほど時間がかかります。時間短縮のために8エポックだけ学習させたい場合は次のようにします。

./scripts/train.sh trainer.max_epochs=8

学習中の GPU の使用状況を確認するには、別のターミナルを開いて rocm-smi コマンドを実行します。次のように出力されます。

$ rocm-smi


======================================= ROCm System Management Interface =======================================
================================================= Concise Info =================================================
Device  [Model : Revision]    Temp    Power   Partitions      SCLK     MCLK     Fan  Perf  PwrCap  VRAM%  GPU%
        Name (20 chars)       (Edge)  (Avg)   (Mem, Compute)
================================================================================================================
0       [0x0c34 : 0x02]       51.0°C  222.0W  N/A, N/A        1675Mhz  1600Mhz  0%   auto  300.0W   12%   100%
        Instinct MI210
================================================================================================================
============================================= End of ROCm SMI Log ==============================================

学習の進み具合は TensorBoard を使用して確認できます。TensorBoard を起動するには別のターミナルを開いて次のようにします。

cd /scratch/$USER/vai-tutorial
source .venv/bin/activate
tensorboard --logdir ./logs/train/runs/

リモートデスクトップで as006 に接続し、Firefox または Chrome で http://localhost:6006 を開くと TensorBoard の画面にアクセスできます。学習が進むにつれロスが下がり精度が上がっていることが分かります。

200エポック学習させたときの精度は95.1%となりました。

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/acc          │    0.9510999917984009     │
│         test/loss         │    0.1570771187543869     │
└───────────────────────────┴───────────────────────────┘

Vitis AI による量子化、コンパイル

Vitis AI による量子化は GPU により高速に行うことができます。上記のチュートリアルで学習したモデルを量子化、コンパイルしてみましょう。学習済みのモデルは ./logs ディレクトリの下に保存されています。

./logs/train/runs/2024-04-10_18-01-47/checkpoints/
├── epoch_196.ckpt
└── last.ckpt

この例ではテスト用のデータセットに対する精度が最も高かったモデルが epoch_196.ckpt として保存されています。以降の手順ではこのファイルを対象に量子化を行います。

量子化の手順

Vitis AI の PTQ (Post Training Quantization) による量子化の手順は次のとおりです。

  1. キャリブレーション
    • モデルにデータを入力し、各層の出力の分布を調べることで量子化係数を決定します
  2. テスト
    • 量子化後のモデルの精度を確認します
  3. デプロイ
    • 量子化したモデルを xmodel 形式でエクスポートします
    • xmodel ファイルはコンパイラへの入力になります

Vitis AI による量子化は正式には Docker 環境で行いますが、公式の Docker イメージにインストールされている PyTorch のバージョンが古いため、ここでは Docker は使用せず量子化に必要なツールをソースコードからビルドして使用します。次のコマンドでビルドを実行します。

cd /scratch/$USER/vai-tutorial
source .venv/bin/activate
./scripts/build_quantizer.sh

キャリブレーションを実行します。ckpt_path= のパスはご自身の環境に合わせて変更してください。

./scripts/quant_calib.sh ckpt_path=./logs/train/runs/2024-04-10_18-01-47/checkpoints/epoch_196.ckpt

量子化に関する情報は quantized ディレクトリ以下に保存されます。

次に量子化後のモデルの精度をテストします。

./scripts/quant_test.sh ckpt_path=./logs/train/runs/2024-04-10_18-01-47/checkpoints/epoch_196.ckpt

量子化前の95.1%からわずかに下がり94.7%となりました。問題ないためこのまま進めます。

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/acc          │    0.9473000168800354     │
│         test/loss         │    0.16378344595432281    │
└───────────────────────────┴───────────────────────────┘

最後にモデルを xmodel 形式でエクスポートします。

./scripts/quant_deploy.sh ckpt_path=./logs/train/runs/2024-04-10_18-01-47/checkpoints/epoch_196.ckpt

./quantized/ResNet9_int.xmodel に量子化したモデルが保存されます。

コンパイルの手順

Vitis AI のコンパイラはオープンソースではないため Docker 環境から使用する必要があります。as006 および as005 では Rootless Docker をセットアップすることでユーザーが Docker を利用できます。

Rootless Docker のセットアップ方法

初めて Rootless Docker を使用する際に以下を実施する必要があります。

1. Rootless Docker の設定ファイルを作成します。

mkdir -p ~/.config/docker
echo "{ \"data-root\" : \"/docker-data/$(id -un)\" }" > ~/.config/docker/daemon.json

2. Rootless Docker をインストールします。

dockerd-rootless-setuptool.sh install

上記を実施しても Docker デーモンに接続できない場合は次のようにしてサービスの再起動を行ってみてください。

systemctl --user restart docker.service

Docker イメージのロード

Vitis AI の Docker イメージをロードします。モデルのコンパイルには GPU は使用しないため CPU のみのイメージを使用します。これには2分ほど時間がかかります。

docker load -i /tools/docker/xilinx/vitis-ai-pytorch-cpu/ubuntu2004-3.5.0.306.tar

docker images コマンドでイメージがロードされたことが確認できます。

$ docker images
REPOSITORY                    TAG                    IMAGE ID       CREATED        SIZE
xilinx/vitis-ai-pytorch-cpu   ubuntu2004-3.5.0.306   10667a5b7322   9 months ago   11.7GB

Docker コンテナを立ち上げるには次のスクリプトを実行します。

./docker/run.sh

うまく立ち上がると次のように表示されます。

==========================================
 
__      ___ _   _                   _____
\ \    / (_) | (_)            /\   |_   _|
 \ \  / / _| |_ _ ___ ______ /  \    | |
  \ \/ / | | __| / __|______/ /\ \   | |
   \  /  | | |_| \__ \     / ____ \ _| |_
    \/   |_|\__|_|___/    /_/    \_\_____|
 
==========================================

Docker Image Version: ubuntu2004-3.5.0.306   (CPU) 
Vitis AI Git Hash: 6a9757a 
Build Date: 2023-06-26
WorkFlow: pytorch

root@as006:/workspace#

xmodel 形式のモデルのコンパイルには vai_c_xir コマンドを使用します。VCK5000 をターゲットにコンパイルするには次のようにします。

vai_c_xir \
  --xmodel ./quantized/ResNet9_int.xmodel \
  --arch /opt/vitis_ai/compiler/arch/DPUCVDX8H/VCK50008PE/arch.json \
  --output_dir ./compiled \
  --net_name ResNet9_Cifar10

VCK5000 向けの DPU 実装は PE のインスタンス数やサポートレイヤーの違いから以下のバリエーションがあります。使用したい DPU 実装に合わせて上記コマンドの --arch オプションへ渡すパスを変更してください。

DPU アーキPE 数タイプarch.json のパス
DPUCVDX8H4miscdwc/opt/vitis_ai/compiler/arch/DPUCVDX8H/VCK50004PE/arch.json
DPUCVDX8H6dwc/opt/vitis_ai/compiler/arch/DPUCVDX8H/VCK50006PEDWC/arch.json
DPUCVDX8H6misc/opt/vitis_ai/compiler/arch/DPUCVDX8H/VCK50006PEMISC/arch.json
DPUCVDX8H8normal/opt/vitis_ai/compiler/arch/DPUCVDX8H/VCK50008PE/arch.json

特に問題がなければ ./compiled ディレクトリ以下にコンパイル後のモデルが出力されます。Docker コンテナは閉じてしまって構いません。

コンパイルされたモデルがどのように DPU に割り付けられたかを Netron を使って確認します。ホスト側のターミナルに戻り、次のコマンドを実行します。

cd /scratch/$USER/vai-tutorial
source .venv/bin/activate
netron ./compiled/ResNet9_Cifar10.xmodel

リモートデスクトップのブラウザで http://localhost:8080 を開きます。

upload と download はホストとデバイスの間のデータ転送を意味します。モデルの中に DPU がサポートしていないレイヤーが含まれているとホスト CPU で処理する必要があり、余計なデータ転送が発生してしまいます。今回使用したモデルの元の実装では nn.MaxPool2d(4) が使用されていてこれに DPU が対応していなかったため、連続するふたつの nn.MaxPool2d(2) に置き換えることで対策しました。

Vitis AI によるモデルの実行

Vitis AI で量子化、コンパイルしたモデルを VCK5000 で動作させてみましょう。ここからは as005 にログインして作業を進めます。

as005 のスクラッチ領域にチュートリアルのリポジトリをコピーします。

cd /scratch/$USER
cp -a /tools/repo/acri-room/vai-tutorial .
cd vai-tutorial

venv 環境を作成し、PyTorch をインストールします。as005 では GPU は使用できないため CPU 版のPyTorch を指定します。

python -m venv .venv
source .venv/bin/activate
pip install --find-links=/tools/pip/pytorch torch==2.2.2+cpu torchvision==0.17.2+cpu torchaudio==2.2.2+cpu

他に必要となるパッケージをインストールします。

pip install -r requirements.txt

モデルの実行に必要となる Vitis AI Runtime 関連のソフトウェアをビルドします。

./scripts/build_runtime.sh

ターゲットの DPU 実装を指定するための環境設定を行います。

 source ./docker/setup_vck5000.sh DPUCVDX8H_8pe_normal

最後の引数にはコンパイル時に指定した DPU 実装を次の中から指定します。

  • DPUCVDX8H_4pe_miscdwc
  • DPUCVDX8H_6pe_dwc
  • DPUCVDX8H_6pe_misc
  • DPUCVDX8H_8pe_normal

as006 でコンパイルしたモデルを as005 にコピーします。予約不要の as106 が as006 とスクラッチ領域を共有しているためこれを利用します。

scp -r as106:/scratch/$USER/vai-tutorial/compiled .

入力データとして CIFAR10 データセットを展開します。

mkdir data
tar xf /tools/data/cifar-10-python.tar.gz -C data

以上でモデルを実行する準備が整いました。

次のように xmodel= にコンパイルしたモデルのパスを渡して実行します。初回は少し時間がかかります。

./scripts/run.sh xmodel=./compiled/ResNet9_Cifar10.xmodel

成功すると次のように表示されます。ターミナルによって画像がうまく表示されない場合はリモートデスクトップで接続して Ubuntu のターミナルから実行してみてください。

より多くの結果を確認したい場合は num_batches= オプションを指定します。

./scripts/run.sh xmodel=./compiled/ResNet9_Cifar10.xmodel num_batches=100

タイトルとURLをコピーしました