MXNet Model Server で ONNXモデルをサーブしてみる

awslabs/mxnet-model-serverのv0.2.0から、ONNXモデルをサーブできるようになったらしいので試してみる。
参考: AWS Machine Learning Blog - Model Server for Apache MXNet introduces ONNX support and Amazon CloudWatch integration

環境構築

クイックスタートを参考にubuntu:16.04なDocker上で環境構築。以下のDockerfileでDocker Imageを作成する。

RUN apt-get update
RUN apt-get install protobuf-compiler libprotoc-dev python-pip build-essential -y
RUN pip install mxnet-model-server

Docker buildしてモデルのサーブをしてみる。mxnetな.modelファイルでは簡単に動かす事ができた。

$ docker build -t mms .
$ docker run -it mms bash
>mxnet-model-server \
  --models squeezenet=https://s3.amazonaws.com/model-server/models/squeezenet_v1.1/squeezenet_v1.1.model

ONNXモデルをサーブする

MXNet model serverで読み込める形式にする

Documentを参考にONNXモデルをMMSでサーブできる形式に変換する。

$ mkdir onnx-squeezenet && cd onnx-squeezenet
$ curl -O https://s3.amazonaws.com/model-server/models/onnx-squeezenet/squeezenet.onnx
$ curl -O https://s3.amazonaws.com/model-server/models/onnx-squeezenet/signature.json
$ curl -O https://s3.amazonaws.com/model-server/models/onnx-squeezenet/synset.txt

$ mxnet-model-export --model-name onnx-squeezenet --model-path .
$ ls -lh onnx-squeezenet.model
-rw-r--r-- 1 root root 4.8M Feb 11 09:38 onnx-squeezenet.model

モデルのサーブをしてみる

できた

$ mxnet-model-server --models squeezenet=onnx-squeezenet.model
[INFO 2018-02-11 09:43:38,738 PID:653 /usr/local/lib/python2.7/dist-packages/mms/mxnet_model_server.py:__init__:93] Initialized model serving.
[WARNING 2018-02-11 09:43:38,742 PID:653 /usr/local/lib/python2.7/dist-packages/mms/model_service/mxnet_model_service.py:__init__:95] Failed to parse epoch from param file, setting epoch to 0
[INFO 2018-02-11 09:43:38,754 PID:653 /usr/local/lib/python2.7/dist-packages/mms/serving_frontend.py:add_endpoint:182] Adding endpoint: squeezenet_predict to Flask
[INFO 2018-02-11 09:43:38,755 PID:653 /usr/local/lib/python2.7/dist-packages/mms/serving_frontend.py:add_endpoint:182] Adding endpoint: ping to Flask
[INFO 2018-02-11 09:43:38,755 PID:653 /usr/local/lib/python2.7/dist-packages/mms/serving_frontend.py:add_endpoint:182] Adding endpoint: api-description to Flask
[INFO 2018-02-11 09:43:38,779 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric CPUUtilization for last 60 seconds is 56.700000
[INFO 2018-02-11 09:43:38,780 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric MemoryUsed for last 60 seconds is 173.000000
[INFO 2018-02-11 09:43:38,785 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric MemoryAvailable for last 60 seconds is 1472.000000
[INFO 2018-02-11 09:43:38,786 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric MemoryUtilization for last 60 seconds is 26.300000
[INFO 2018-02-11 09:43:38,787 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric DiskUsed for last 60 seconds is 30.000000
[INFO 2018-02-11 09:43:38,789 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric DiskAvailable for last 60 seconds is 29.000000
[INFO 2018-02-11 09:43:38,789 PID:653 /usr/local/lib/python2.7/dist-packages/mms/metric.py:start_recording:214] Metric DiskUtilization for last 60 seconds is 50.400000
[INFO 2018-02-11 09:43:38,790 PID:653 /usr/local/lib/python2.7/dist-packages/mms/mxnet_model_server.py:start_model_serving:107] Service started successfully.
[INFO 2018-02-11 09:43:38,790 PID:653 /usr/local/lib/python2.7/dist-packages/mms/mxnet_model_server.py:start_model_serving:108] Service description endpoint: 127.0.0.1:8080/api-description
[INFO 2018-02-11 09:43:38,790 PID:653 /usr/local/lib/python2.7/dist-packages/mms/mxnet_model_server.py:start_model_serving:109] Service health endpoint: 127.0.0.1:8080/ping
[INFO 2018-02-11 09:43:38,792 PID:653 /usr/local/lib/python2.7/dist-packages/werkzeug/_internal.py:_log:88]  * Running on http://127.0.0.1:8080/ (Press CTRL+C to quit)

変換後のモデルの中身

変換後のモデルの中身を見てみる。

root@05cd6f4910d9:/onnx-squeezenet# ls -l onnx-squeezenet
total 4904
-rw-r--r-- 1 root root     536 Feb 11 09:39 MANIFEST.json
-rw-r--r-- 1 root root    1809 Feb 11 09:39 mxnet_vision_service.py
-rw-r--r-- 1 root root    1836 Feb 11 09:39 mxnet_vision_service.pyc
-rw-r--r-- 1 root root     259 Feb 11 09:39 signature.json
-rw-r--r-- 1 root root 4945334 Feb 11 09:39 squeezenet-0000.params
-rw-r--r-- 1 root root   21696 Feb 11 09:39 squeezenet-symbol.json
-rw-r--r-- 1 root root   31675 Feb 11 09:39 synset.txt

mxnetな形式に変換されてた。サーバ起動時に.onnxが読み込まれるわけではなく、事前に変換する方式なんですね。