ONNX形式モデルをDocker/CPU上のTensorFlow Servingで予測する

抄訳

• ONNXはNCHW形式のデータを受け付ける
• 他方通常ビルド形式のTensorflowはNHWC形式のデータしか受け付けない
• 実際のサーバではGPUを使いたくない
• MKLを使用するとCPUでもHCHW形式の予測が可能になる
• tensorflow-serving公式のDockerfileはMKL形式でビルドしていない
• 公式Dockerfileを改造したら予測が動くようになった

詳細

Pytorchで訓練したモデルをONNX形式にexportしてtensorflow-servingで配信というフローがやりたかった。

公式のDocker Imageを普通にpullすると、config=nativeoptでビルドされたtensorflowを使用することになる為、MKLを導入するにはビルドの仕方を改造する必要がある。

公式のDockerfileからローカルでビルドを行う方法についてはこちらを参照

custom_build_procedure

# just clone public repository
git clone https://github.com/tensorflow/serving
cd serving

# edit Dockerfiles
sed -i 's#ubuntu:16.04#triage/python3.6-mkl#g' tensorflow_serving/tools/docker/Dockerfile.devel
sed -i 's#ubuntu:16.04#triage/python3.6-mkl#g' tensorflow_serving/tools/docker/Dockerfile

# build with custom options
docker build --pull -t $USER/tensorflow-serving-devel --build-arg TF_SERVING_BUILD_OPTIONS="--config=nativeopt --config=mkl" -f t
ensorflow_serving/tools/docker/Dockerfile.devel .
docker build -t $USER/tensorflow-serving --build-arg TF_SERVING_BUILD_IMAGE=$USER/tensorflow-serving-devel -f tensorflow_serving/tools/docker/Dockerfile .

# run
export MODEL_NAME="your model name"
export MODEL_DIR="where your model is"
docker run --rm -p 8501:8501 --name $MODEL_NAME --mount type=bind,source=$MODEL_DIR -e MODEL_NAME=$MODEL_NAME -t $USER/tensorflow-serving:latest

解説

基本は公式ガイダンスに従っています。

• MKL-DNNのライブラリが動作時に必要なので、ベースimageを ubuntu:16.04 から triage/python3.6-mklに変更
• ビルドオプションに--config=mklを追加

あまり劇的に魔改造してるわけじゃないから、tensorflow-servingのバージョンが変わってDockerfileが編集されてもこのまま動作してくれるでしょう...