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

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

![NNVM Archtecture](http://www.tvmlang.org/images/nnvm/nnvm_compiler_stack.png)

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


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

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

NNVMの公式チュートリアルを参考にして進めていきます。
[Compile Keras Models](http://nnvm.tvmlang.org/tutorials/from_keras.html#sphx-glr-tutorials-from-keras-py)

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


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

## 必要環境の準備

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

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

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

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

keras環境を導入します。

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


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

```bash
$pip install pillow
```

## NNVMのインストール

[公式ドキュメント](http://nnvm.tvmlang.org/how_to/install.html)を参考に導入していきます。

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

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

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

ビルドを実行します。

```bash
$cd /path/to/nnvm
$make
```

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

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

### Python Packageのインストール

pythonの環境パスを通します

```bash
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を修正します。

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

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

```bash:config.mk
LLVM_CONFIG=/usr/lib/llvm-5.0/bin/llvm-config
```


ビルドを実行します。

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

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

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

### Python Packageのインストール

pythonの環境パスを通します

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

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

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

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

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

```python
import nnvm
import tvm
```


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

[公式チュートリアル](http://nnvm.tvmlang.org/tutorials/from_keras.html#sphx-glr-tutorials-from-keras-py)を少しだけ改変し、NNVM/TVMによりKerasのResNet50モデルをCPU向けにコンパイルするサンプルを実行します。

```python: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]))
```

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

```bash: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](http://www.tvmlang.org/2017/10/06/nnvm-compiler-announcement.html)
    - [NNVM Documentation](http://nnvm.tvmlang.org/)
    - [TVM Documentation](http://docs.tvmlang.org/index.html)
    - [Introducing NNVM Compiler](https://aws.amazon.com/jp/blogs/machine-learning/introducing-nnvm-compiler-a-new-open-end-to-end-compiler-for-ai-frameworks/)

- ソースコード (GitHub)
    - [Distributed (Deep) Machine Learning Community / NNVM](https://github.com/dmlc/nnvm)