はじめに
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を使っているあたりが萌えます。
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じゃないのかなあ。。
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"
と言われエラーになります。まだ追えていません。
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))
デプロイファイル一式を使って推論を行うのが下記です。ミニバッチ１にしたのですがこのあたりは保存時に合わせるのでしょうね。
#!/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を使って推論をするコードも書きました。ミニバッチ１個ずつ演算しているのは比較のためです。
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になりましたので下記に書いていたインストール手順は消します。