この記事について
推論コンパイラNNVM/TVMによって、推論時に利用できるライブラリモジュールを生成する方法について記載します。ライブラリ化することにより、推論用のターゲットデバイスに持っていきやすくなります。
この記事の目標
- NNVMでKeras付属のモデルをCPUターゲット向けにコンパイルします
- NNVMでコンパイルしたモジュールをライブラリ化します
- TVMでライブラリ化したモジュールを呼び出し、推論を実行します
ネットワークモデルのライブラリ化
NNVMを使ってモデルをライブラリ化していきます。モデル本体はKeras付属のResNet50を用います。この例では、推論側に持っていきやすいように、入力画像を前処理した後のデータをpickleでシリアライズ化していますが、必須ではありません。
import nnvm
import keras
import numpy as np
import pickle
def download(url, path, overwrite=False):
import os
if os.path.isfile(path) and not overwrite:
print('File {} exists, skip.'.format(path))
return
print('Downloading from url {} to {}'.format(url, path))
try:
import urllib.request
urllib.request.urlretrieve(url, path)
except:
import urllib
urllib.urlretrieve(url, path)
# ----------------------------------------
# ResNet50を取得します
# ----------------------------------------
weights_url = ''.join(['https://github.com/fchollet/deep-learning-models/releases/',
'download/v0.2/resnet50_weights_tf_dim_ordering_tf_kernels.h5'])
weights_file = 'resnet50_weights.h5'
download(weights_url, weights_file)
keras_resnet50 = keras.applications.resnet50.ResNet50(include_top=True, weights=None,
input_shape=(224,224,3), classes=1000)
keras_resnet50.load_weights('resnet50_weights.h5')
# -----------------------------------------
# 入力画像を取得し、
# 前処理済みのデータを生成します
# -----------------------------------------
from PIL import Image
from keras.applications.resnet50 import preprocess_input
img_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
download(img_url, 'cat.jpg')
img = Image.open('cat.jpg').resize((224, 224))
# input preprocess
data = np.array(img)[np.newaxis, :].astype('float32')
data = preprocess_input(data).transpose([0, 3, 1, 2])
print('data', data.shape)
# Pickleにより推論入力に使う入力データをシリアライズ化
with open('data.pickle', 'wb') as data_file:
pickle.dump(data , data_file)
# -----------------------------------------
# モデルをNNVMによりコンパイルし
# 関数ライブラリを出力します
# -----------------------------------------
# convert the keras model(NHWC layout) to NNVM format(NCHW layout).
sym, params = nnvm.frontend.from_keras(keras_resnet50)
# compile the model
target = 'llvm'
shape_dict = {'data': data.shape}
with nnvm.compiler.build_config(opt_level=2):
graph, lib, params = nnvm.compiler.build(sym, target, shape_dict, params=params)
# ライブラリをエクスポート
lib.export_library('libgraph.so')
# JSONファイルを出力
with open('graph.json',"w") as fo:
fo.write(graph.json())
# パラメータを出力
with open('graph.params', "w") as fo:
fo.write(nnvm.compiler.save_param_dict(params))
実行すると、以下のようにライブラリ等が生成されます。
$python export.py
$ls
data.pickle graph.json graph.params libgraph.so
TVMによる呼び出しコード
生成した以下のファイルを読み込んで推論を行うプログラムを作成します。
- libgraph.so (関数ライブラリ)
- graph.json (モデル構成)
- graph.param (モデルパラメータ)
- data.pickle (前処理済み入力データ)
import pickle
import tvm
import numpy as np
from tvm.contrib import graph_runtime
# ----------------------------------------
# pickleによりシリアライズ化した
# 入力データを読み出します
# ----------------------------------------
with open('data.pickle', 'rb') as data_file:
data = pickle.load(data_file)
print("image data loaded")
# -----------------------------------------
# ライブラリを読み込みます
# -----------------------------------------
loaded_lib = tvm.module.load('./libgraph.so')
print("lib loaded")
# -----------------------------------------
# 各種パラメータを読み込みます
# -----------------------------------------
loaded_json = open('graph.json').read()
print("json loaded")
loaded_params = bytearray( open('graph.params', "rb").read())
print("params loaded")
module = graph_runtime.create(loaded_json, loaded_lib, tvm.cpu(0))
module.load_params(loaded_params)
# -----------------------------------------
# 実行します
# -----------------------------------------
# set input
module.set_input('data', tvm.nd.array(data.astype('float32')))
# m.set_input(**params)
# execute
print("start prediction")
module.run()
print("end prediction")
# get outputs
out_shape = (1000,)
tvm_out = module.get_output(0, tvm.nd.empty(out_shape, 'float32')).asnumpy()
top1_tvm = np.argmax(tvm_out)
# -----------------------------------------
# 結果を表示します
# -----------------------------------------
synset_name = 'synset.txt'
with open(synset_name) as f:
synset = eval(f.read())
print('NNVM top-1 id: {}, class name: {}'.format(top1_tvm, synset[top1_tvm]))
実行すると、以下のように推論結果が表示されます。
$python predict.py
image data loaded
lib loaded
json loaded
params loaded
start prediction
end prediction
NNVM top-1 id: 282, class name: tiger cat
補足
今回の例では、コンパイルも推論もx86上で行っています。コンパイルはホスト上、推論はARMなどのターゲットボード上で行う場合は、NNVM runtimeとエクスポートするライブラリをターゲット向けにコンパイルする必要があります。
targetで指定するLLVMのパラメータ引数の変更
# ARMv7の例
target = 'llvm -mtriple=armv7l-none-linux-gnueabihf -mcpu=cortex-a53 -mattr=+neon'
# ARMv8の例
target = 'llvm -mtriple=aarch64-none-linux-gnu -mcpu=cortex-a57 -mattr=+neon'
# その他は変更なしで実行できます
shape_dict = {'data': data.shape}
with nnvm.compiler.build_config(opt_level=2):
graph, lib, params = nnvm.compiler.build(sym, target, shape_dict, params=params)
TVM(Runtime)のクロスコンパイル
Makefileにクロスコンパイラの情報を追加します。(CMakeで変更することもできます)
/usr/bin、/usr/aarch64-linux-gnuにARMコンパイラをインストールしている場合
# The flags
CC=/usr/bin/aarch64-linux-gnu-gcc
CXX=/usr/bin/aarch64-linux-gnu-g++
LDFLAGS = -L /usr/aarch64-linux-gnu/lib -pthread -lm -ldl
ビルドを実行します。
$cd /path/to/tvm
$make
$ls lib
libtvm_runtime.so