NNVM/TVMでKerasのモデルをコンパイルする

この記事について

NNVMはDeep learningの推論用コンパイラの一つで、AWSなどによりサポートされています。NNVMのアーキテクチャは以下の通りです。(公式HPより)

この記事では、NNVM/TVMの導入と簡単なモデルのコンパイルについて記載します。

この記事の目標

今回は上図で示されているところの、Keras->NNVM->TVM->LLVM->x86の手順を実行します。

• NNVM/TVMをUbuntu(x86)上の環境に導入します
• KerasのサンプルモデルをNNVMでコンパイルします
• コンパイルされたNNVMモデルからTVMでCPU向けの推論モジュールに変換します
• NNVM/TVMで作成したモジュールで推論を実行します

NNVMの公式チュートリアルを参考にして進めていきます。
Compile Keras Models

Ubuntu 16.04、g++ 5.4.0、python 2.7.12の環境で構築します。

NNVM/TVMの導入

NNVM/TVMはその名の通り、大きくNNVMモジュールとTVMモジュールの二つのモジュールが協調して動作します。NNVMが各種フレームワークで学習されたモデルの最適化を担当し、TVMが最適化された学習モデルを実機動作可能なプログラムに変換するバックエンドを担当します。

必要環境の準備

python開発環境を導入します。

$sudo apt-get update
$sudo apt-get install -y python python-dev python-setuptools gcc libtinfo-dev zlib1g-dev

LLVM-5.0を導入します。ソースコードから入手してもよいですが、プリビルドモジュールがあるため、そちらから入手します。プリビルドモジュールの取得方法はここに記載されています。準備ができたらaptによりインストールします。

$sudo apt-get update
$sudo apt-get install llvm-5.0

keras環境を導入します。

$pip install -U keras --user
$pip install -U tensorflow --user 

チュートリアルで必要となるライブラリ(PIL)を導入します。

$pip install pillow

NNVMのインストール

公式ドキュメントを参考に導入していきます。

NNVM Shared Library (libnnvm_compiler.so)のビルド

ソースコードを取得します。

$git clone --recursive https://github.com/dmlc/nnvm
$cd nnvm
$git submodule init
$git submodule update --recursive

ビルドを実行します。

$cd /path/to/nnvm
$make

成功すると、libnnvm_compiler.soが生成されます。

$cd /path/to/nnvm/lib
$ls
libnnvm.a  libnnvm_compiler.so

Python Packageのインストール

pythonの環境パスを通します

export PYTHONPATH=/path/to/nnvm/python:${PYTHONPATH}

nnvm packageをカレントユーザー用にインストールします

cd /path/to/nnvm/python
python setup.py install --user

TVMのインストール

TVM Shared Libraryのビルド

config.mkを修正します。

$cd /path/to/nnvm/tvm/make
$vi config.mk

LLVMの変数にllvm-configのパスを設定します。(/usr/lib/llvm-5.0にインストールした場合)

config.mk

LLVM_CONFIG=/usr/lib/llvm-5.0/bin/llvm-config

ビルドを実行します。

$cd /path/to/nnvm/tvm
$make

成功すると、libtvm_runtime.so libtvm.so libtvm_topi.soが生成されます。

$cd /path/to/nnvm/tvm/lib
$ls
libtvm_runtime.so  libtvm.so  libtvm_topi.so

Python Packageのインストール

pythonの環境パスを通します

$export PYTHONPATH=/path/to/tvm/python:/path/to/tvm/topi/python:${PYTHONPATH}

tvm packageをカレントユーザー用にインストールします。

cd /path/to/tvm/python
python setup.py install --user

cd /path/to/tvm/topi/python
python setup.py install --user

これで利用準備が完了し、NNVM/TVMをimportできるようになりました。

import nnvm
import tvm

サンプルプログラムの実行

公式チュートリアルを少しだけ改変し、NNVM/TVMによりKerasのResNet50モデルをCPU向けにコンパイルするサンプルを実行します。

keras_example.py

import nnvm
import tvm
import keras
import numpy as np

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のWeightをダウンロード
#-------------------------------
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 matplotlib import pyplot as plt
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))
##plt.imshow(img)
##plt.show()
# input preprocess
data = np.array(img)[np.newaxis, :].astype('float32')
data = preprocess_input(data).transpose([0, 3, 1, 2])
print('data', data.shape)


#-------------------------------------
# NNVMによるモデルのコンパイル
# (targetをCPU(LLVM)に変更しています)
#-------------------------------------
# convert the keras model(NHWC layout) to NNVM format(NCHW layout).
sym, params = nnvm.frontend.from_keras(keras_resnet50)
# compile the model
##target = 'cuda'
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)


#-------------------------------------
# TVMの実行
# (コンテキストをCPUに変更しています)
#-------------------------------------
from tvm.contrib import graph_runtime
##ctx = tvm.gpu(0)
ctx = tvm.cpu(0)
m = graph_runtime.create(graph, lib, ctx)
# set inputs
m.set_input('data', tvm.nd.array(data.astype('float32')))
m.set_input(**params)
# execute
m.run()
# get outputs
out_shape = (1000,)
tvm_out = m.get_output(0, tvm.nd.empty(out_shape, 'float32')).asnumpy()
top1_tvm = np.argmax(tvm_out)


#-------------------------------------
# 比較実行
#-------------------------------------
synset_url = ''.join(['https://gist.githubusercontent.com/zhreshold/',
                      '4d0b62f3d01426887599d4f7ede23ee5/raw/',
                      '596b27d23537e5a1b5751d2b0481ef172f58b539/',
                      'imagenet1000_clsid_to_human.txt'])
synset_name = 'synset.txt'
download(synset_url, synset_name)
with open(synset_name) as f:
    synset = eval(f.read())
print('NNVM top-1 id: {}, class name: {}'.format(top1_tvm, synset[top1_tvm]))
# confirm correctness with keras output
keras_out = keras_resnet50.predict(data.transpose([0, 2, 3, 1]))
top1_keras = np.argmax(keras_out)
print('Keras top-1 id: {}, class name: {}'.format(top1_keras, synset[top1_keras]))

成功すると、以下のように推論結果が表示されます。

output

File synset.txt exists, skip.
NNVM top-1 id: 282, class name: tiger cat
Keras top-1 id: 282, class name: tiger cat

関連情報

• ドキュメント

  • 公式HP
  • NNVM Documentation
  • TVM Documentation
  • Introducing NNVM Compiler

• ソースコード (GitHub)

  • Distributed (Deep) Machine Learning Community / NNVM