Kria Starter Kitで、Vitis AIを動かす
カスタムモデルを使用してみる

公開日：2023年2月22日

• 日本語
• 英語

MNISTの概要

MNISTは0-9の手書き数字の画像分類です。

TensorFlowでデータセットが簡単に使用できるため、AIモデルの作り方、学習済みAIモデルの保存、使用方法等の学習に適しています。

ボード向けコンパイル方法の概要

以下がコンパイルのシーケンス概要です。

1. 学習済みAIモデルを用意
2. AIモデルを量子化
3. ボード用にコンパイル

2,3の機能がVitis-AIで提供されており、学習済みAIモデルがあれば、ボード向けにコンパイルして実行することができます。

量子化により、AI推論の精度が変わりますので、量子化前後のAI推論結果は完全一致しない点に注意してください。

3.ではボードのDPUアーキテクチャに合わせてコンパイルする必要がありますので注意してください。

TensorFlow2によるMNISTの学習、量子化

Vitis-AI開発環境を起動し、conda activate vitis-ai-tensorflow2を入力してTensorFlow2環境を有効にします。

AIモデルの学習時はVitis-AI開発環境は必須ではないですが、TensorFlow2環境が整っているので、本記事ではAIモデルの学習もVitis-AI開発環境で行います。

作業ディレクトリを作成、移動します。

AIモデルの作成、MNISTの学習、量子化して保存する処理を行います。

以下をmy_mnist_test.pyとして保存してpython my_mnist_test.pyとして実行してください。

内容についてはコード中のコメントを参照してください。

AIモデル層の内容や学習パラメータは深く考慮したものではないのでご注意ください。

また、毎回結果が変わらないように乱数固定しています。

import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow_model_optimization.quantization.keras import vitis_quantize

# 乱数固定
tf.random.set_seed(0)

# 学習、テスト用データセット準備
# この関数でMNISTの学習に必要なデータセットを扱うことができます
(x_train, y_train),(x_test, y_test) = tf.keras.datasets.mnist.load_data()

# 学習用画像データを正規化します
x_train = x_train/255.0
x_test = x_test/255.0

# AIモデル層作成
layers_list = []

# 一次元配列に変換
layers_list.append(tf.keras.layers.Flatten(input_shape=(28, 28, 1)))
# 全結合層(128)、activationはReLuを指定
layers_list.append(tf.keras.layers.Dense(128, activation='relu'))
# Dropout(0.2)
layers_list.append(tf.keras.layers.Dropout(0.2))
# 全結合層(10)、activationはsoftmaxを指定、ここが最終的な出力となる
layers_list.append(tf.keras.layers.Dense(10, activation='softmax'))

# AIモデル設定
mnist_model = tf.keras.models.Sequential(layers_list)

# AIモデル学習設定
mnist_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# AIモデル学習実行
mnist_model.fit(x_train, y_train, batch_size=128, epochs=30, validation_data=(x_test, y_test))

# AIモデルの量子化、保存
quantizer = vitis_quantize.VitisQuantizer(mnist_model)
# 渡すデータセット数はラベル無しで100-1000で良いそうなので500としておく
quantized_model = quantizer.quantize_model(calib_dataset=x_train[0:500])
quantized_model.save("quantized_model.h5")

KV260向けコンパイル

続いてVitis-AI環境でAI学習、量子化済みのAIモデルをターゲットボード向けにコンパイルします。

まず、DPUアーキテクチャをB4096で指定するためのarch.jsonを作成します。

本記事では、B4096(fingerprint:0x101000016010407)を指定しますので以下をarch.jsonとして保存してください。

{
    "target": "DPUCZDX8G_ISA1_B4096"
}

コンパイルはvai_c_tensorflow2コマンドを使用します。

以下の内容をmy_mnist_compile.shとして保存し、作業ディレクトリで実行してください

#!/bin/bash
ARCH=./arch.json
COMPILE_KV260=./outputs_kv260

vai_c_tensorflow2 \
    --model      ./quantized_model.h5 \
    --arch       ${ARCH} \
    --output_dir ${COMPILE_KV260} \
    --net_name   my_mnist_test

出力されたmy_mnist_test.xmodelがコンパイル済みAIモデルとなります。

本記事で使用するファイルのダウンロード

KV260でAI推論するためのJupyter Notebook, テスト画像をダウンロードしておきます。

以下からファイルをダウンロードしてください。

ファイルのダウンロード

中身は以下となります。

AI推論用ファイルの配置

• KV260向けコンパイル
• 本記事で使用するファイルのダウンロード

で取得した3ファイルをmicroSDカード上に配置します。

配置方法はModel Zooを使用してみるを参照してください。

コンパイルしたAIモデルの動作確認

コンパイルしたAIモデルを使用する方法は、Model Zooを使用してみると同じです。

Jupyter NotebookでCustomModel_test.ipynbを選択して各セルを実行してください。

AI推論時にエラーとならないことが確認できればokです。

実行結果は以下の画像のようになります。

3の数字のテスト画像を入力としてAI推論を実行、正しく判定できることが確認できます。

まとめ

AIモデルを作成してボード向けコンパイルを行い、KV260で実行してみました。

ユーザが独自に作成したAIモデルを、KV260上で実行できることを確認できました。

但し、使用される全てのオペレータに対応しているわけではないので、AIモデルによっては量子化、コンパイルが失敗、または性能が十分でない可能性があります。

詳細はサポート情報を参照してください。

※文中に記載されている各種名称、会社名、商品名などは各社の商標もしくは登録商標です。

• 前の記事を読む
• 一覧に戻る
• 次の記事を読む