はじめてのNNVM

はじめに

NNVMは各種フロントエンドなディープラーニングフレームワーク(Caffe/Keras/MXNet/PyTorch/Caffe2/CNTK)で構築した演算グラフを、TVMというテンソル演算スタックを介してさまざまなバックエンド(LLVM,OpenCL,Metal,CUDA)を用いたランタイムに変換するコンパイラです。原理的には学習にも使えるんじゃないかと思うのですが、基本的には推論のデプロイ用というかんじですね。

すごく乱暴にいうと、ラズパイ等の非力なデバイスでなるべく速く動くようにするコンパイラ、です。

公式リリースの絵がとてもわかりやすいです。

日本製フレームワーク(Chainer, NNabla)がガン無視されていて悲しいですね。。

残念ながら現時点では、PyTorchからONNX経由での経路は、OSXでは通りませんでした（後述）。素直にサポートされているMXNetで試してみます。
（2017/12/17追記)PyTorchでの変換は別記事を書きました。こちらです。

下記実験に使ったファイルはこちらに置いておきました。また、インストール手順等は記事の後半にまとめてあります。PyTorchやONNXのビルドは今回は無駄でしたがそのうちリベンジしたいと思います。

MXNetでの学習とモデルの保存

毎度恐縮ですが、関数近似で学ぶ chainer とディープラーニングの例でやります。

まずネットを作成します。MXNet初体験だったのですが、gluonというパッケージをつかうとChainer風に書けるようですね。引数にFを使っているあたりが萌えます。

net.py

from mxnet import gluon
from mxnet.gluon import nn

class Net(gluon.HybridBlock):
    def __init__(self, **kwargs):
        super(Net, self).__init__(**kwargs)
        with self.name_scope():
            self.fc1 = nn.Dense(64)
            self.fc2 = nn.Dense(256)
            self.fc3 = nn.Dense(1)
    def hybrid_forward(self, F, x):
        h = F.relu(self.fc1(x))
        h = F.relu(self.fc2(h))
        y = self.fc3(h)
        return y

いつもはLeakyReLUを使うのですが、NNVMでの変換に失敗したのでReLUにしました。ReLUだと初期値次第で収束したりしなかったりする経験があるのですが、少しネットを大きめにすると学習できました。

学習は下記のようにすればできました。L2Lossをbatch_axis=1にしたのですが、これ0じゃないのかなあ。。

train.py

import mxnet as mx
from mxnet import autograd
import logging
logging.getLogger().setLevel(logging.DEBUG)

import numpy as np
from net import *

model=Net()
model.collect_params().initialize(mx.init.Xavier(), ctx=mx.cpu())

def get_batch(n):
    x = np.random.random(n)
    y = np.exp(x)
    return x,y

trainer = gluon.Trainer(model.collect_params(), 'adam')

for i in range(10000):
    with autograd.record():
        x,y = get_batch(100)
        data = mx.nd.array(x).reshape((100,1))
        label = mx.nd.array(y).reshape((100,1))
        output = model(data)
        L = gluon.loss.L2Loss(batch_axis=1) #?!
        loss = L(output, label)
        print loss.asnumpy()
        loss.backward()

    trainer.step(data.shape[0])

model.save_params('model')

保存されたmodelは gluon.HybridBlock というクラスですが、これを直接NNVMに食わせることができます。gluon model-zooにさまざまな著名モデルがあるのでこれらも使うことができるでしょう。

NNVM内部表現への変換はnnvm.frontend.from_mxnet()という関数を使います。ONNXではfrom_onnx()を使えるようですが所々制限があるようです。

下記のようにすると、deploy.dylib, deploy.json, deploy.params というデプロイ用のファイル一式が生成されます。

targetにllvmを指定していますが、ここにcudaとかopenclとかを書けば他のバックエンドを使えます。コンテキストのtvm.cpu(0) はtvm.gpu(0)などに差し替えればGPU版になります。

（2017/12/17追記)NNVMをアップデートすると下記が通らなくなりました。

"only support 2-dim dense"

と言われエラーになります。まだ追えていません。

compile.py

import mxnet as mx
import numpy as np

from net import *

model=Net()
model.load_params("model",mx.cpu(0))

import nnvm
import nnvm.compiler
import tvm
from tvm.contrib import graph_runtime, util

sym, params = nnvm.frontend.from_mxnet(model)

target = 'llvm'
shape_dict = {'data': (1,1)}
graph, lib, params = nnvm.compiler.build(sym, target, shape_dict, params=params, dtype="float32")
module = graph_runtime.create(graph, lib, tvm.cpu(0))

lib.export_library("deploy.dylib")
with open("deploy.json", "w") as fo:
    fo.write(graph.json())
with open("deploy.params", "wb") as fo:
    fo.write(nnvm.compiler.save_param_dict(params))

デプロイファイル一式を使って推論を行うのが下記です。ミニバッチ１にしたのですがこのあたりは保存時に合わせるのでしょうね。

replay_nnvm.py

#!/usr/bin/env python

import numpy as np
import nnvm.compiler
import tvm
from tvm.contrib import graph_runtime, util

from matplotlib import pyplot as plt

import time

loaded_lib = tvm.module.load("deploy.dylib")
loaded_json = open("deploy.json").read()
loaded_params = bytearray(open("deploy.params", "rb").read())

module = graph_runtime.create(loaded_json, loaded_lib, tvm.cpu(0))
params = nnvm.compiler.load_param_dict(loaded_params)
module.load_params(loaded_params)

shape = (1,1)
x_np = np.linspace(0,1,100).astype("float32")

times=[]

outs =np.array([])
start = time.time()
for x in x_np:
    module.run(data=np.array([x]))
    out=module.get_output(0, out=tvm.nd.empty(shape)).asnumpy()[0]
    outs=np.append(outs,out)

elasped_time=time.time() - start
print(elapsed_time)

plt.plot(np.exp(x_np),"b")
plt.hold(True)
plt.plot(outs,"r")
plt.show()

比較用に、MXNetを使って推論をするコードも書きました。ミニバッチ１個ずつ演算しているのは比較のためです。

replay.py

import mxnet as mx
import numpy as np

from matplotlib import pyplot as plt

from net import *

import time

model=Net()
model.load_params("model",mx.cpu(0))

eval_data=np.linspace(0,1,100)

outs=np.array([])
start = time.time()
for x in eval_data:
    out = model(mx.nd.array([x])).asnumpy()
    outs=np.append(outs,out)

elasped_time=time.time() - start
print(elsped_time)

plt.plot(np.exp(eval_data))
plt.plot(outs,"r")

plt.show()

速度比較

100回あたりの演算時間を比較してみました。

Framework	elasped time[sec]
MXNet	0.0533
NNVM(LLVM)	0.0306

お。わりと速いです。まあこの程度のネットだと差が出にくいというか妥当な比較ではないかもしれませんので参考まで。この比較にOpenCL版を並べたかった。。

最後に

まあまだ出始めなのでいろいろと不都合があるのは仕方ないとして、python で学習した結果がそのままダイナミックリンクライブラリにまで落とせるのは大変魅力ですね。クロスコンパイルまでできればラズパイ等々に持っていくのがとても楽になりそうです。

さて、とはいえ、この手の変換ソフトウェアでありがちなことですが、各フレームワークの進化にすべてフル対応するのはなかなか大変なので、変換経路ごとに罠があると思ったほうが良いでしょう（笑）。ただ、標準化に成功すれば（ちゃんと流行すれば）、各フレームワーク側が追従してくれるようになるでしょうから、今後の動向に期待です。

残念ながら現時点では、直接読み込めるフロントエンドはMXNetだけで、それ以外はCoreMLかONNX経由です。MXNet以外はAPIドキュメントもないのでまだこれからという感じでしょうか。ONNXで対応している演算はこちらを見るとわかります。

また残念ながら、PyTorchのONNX exportは現状のリリースでは非対応で、masterをソースからビルドする必要があります。そうやってもOSXではテストが通らず、もちろん変換も通りませんでした。

またまた残念ながら、今のところ対応バックエンドはCUDA,LLVMのみのようです。MacBookのGPUを使いたいのでOpenCLが通ると嬉しかったのですが、現時点では非対応です。ただ、演算部分であるTVMでごく単純なグラフ演算はOpenCLを経由してMacBookのGPU(Intel HD Graphics 4000)で動いたので、それほど遠くない未来に対応できるのではないでしょうか。

期待としては、OSX上でPyTorchでの学習結果をOpenCLで回せるところまで行きたかったのですが、現時点ではいろいろなカベがあって実現できませんでした。

残念がってないでOSSなので貢献して動くようにしろよって話ですね。

さて、これ以降はインストール手順です。

各種インストール

環境は下記です。

• OSX 10.11
• pyenv
• minoconda2-4.0.5

protobufのインストール

手元のprotobufが3.3.0でした。

brew upgrade protobuf

で3.4.1を入れました。ONNXが3.4を要求するので先に入れておきましょう。

NNVMのインストール

LLVMをインストールします。

brew install llvm

bash_profileに下記を追加します。

export PATH=/usr/local/opt/llvm/bin:$PATH

NNVMのレポジトリを落とします。recursiveにすればtvmも同時に落ちてきます。

git clone --recursive https://github.com/dmlc/nnvm

まずtvmをビルドします。

cd nnvm/tvm

CMakeLists.txtを修正して、OpenCLとLLVMを有効にします。

tvm_option(USE_OPENCL "Build with OpenCL" ON)
tvm_option(USE_LLVM "Build with LLVM" ON)

ではビルドします。

mkdir build
cd build
cmake ..
make

build/*.dylib ができました。次にtvm/pythonのビルドとインストールです。

cd ../python
python setup.py install

tvm/topiも必要かもしれません。

cd nnvm/tvm/topi/python
python setup.py install

つぎにnnvmのビルドです。

cd ../../
make

nnvm/pythonのビルドです。

cd python
python setup.py install --user

あとは.bash_profileに下記を書けば終了。

export PYTHONPATH=$HOME/git/nnvm/python:${PYTHONPATH}
export LD_LIBRARY_PATH=$HOME/git/nnvm/tvm/build:${LD_LIBRARY_PATH}

ONNXのインストール

git clone --recursive https://github.com/onnx/onnx.git
cd onnx
pyenv local miniconda2-4.0.5
MACOSX_DEPLOYMENT_TARGET=10.11 CC=clang CXX=clang++ python setup.py install

condaで入れられるのでソースからビルドする必要はなかったかもしれません。

PyTorchのインストール

pipでOkになりましたので下記に書いていたインストール手順は消します。