【備忘録】PyTorchのモデルをonnx形式を経由してOpenVINO用に変換する

2020.09.16

せーのでございます。

コンピュータビジョンを使う際にIntelのOpenVINOというツールキットは足回りを整備してくれていて、非常に使いやすく生産性のあがるものです。システムをIntel系のチップが載ったマシンで動かす場合は必須といっても良いでしょう。

このOpenVINOを使って自作のモデルを動かしたい時にはOpenVINO用の形式にモデルを最適化(optimize)します。

TensorFlowやCaffeなどのフレームワークは直接OpenVINOのオプティマイザーに突っ込んであげれば変換されるのですが、 ChainerやPyTorchで作ったモデルは一旦ONNXと呼ばれるフレームワーク共通のフォーマットに変換してあげて、それをOpenVINOに取り込むことになります。

今回は備忘録の意味も含めましてその手順をご紹介いたします。
今回は画像を骨格検知するために学習したモデルをOpenVINO形式に変換します。

PyTorch => ONNX

最初にPyTorchからONNXへ変換するメソッド全文を載せます。

def convert_to_onnx(net, output_name, single_person, input_size):
    input = torch.randn(1, 3, input_size[0], input_size[1])
    input_layer_names = ['data']
    output_layer_names = ['stage_0_output_1_heatmaps', 'stage_0_output_0_pafs',
                          'stage_1_output_1_heatmaps', 'stage_1_output_0_pafs']

    torch.onnx.export(net, input, output_name, verbose=True, input_names=input_layer_names,
                      output_names=output_layer_names)

PyTorchからONNX形式に変換するメソッドは既にPyTorchのtorchモジュールに入っています。

仕組みとしてはtorchの中にinputとしてテンソルを作ります。

input = torch.randn(1, 3, width, height)

画像は「幅」「高さ」「RGBスケール」「色」という要素でできています。
つまり、上のrandnの引数は左から「RGBか白黒か」「R、G、Bそれぞれの値」「X座標」「Y座標」ということになります。 10 x 10のカラー画像があったとしたら、(1,1)の座標のRの色、Gの色、Bの色、(2,1)の座標のRの色、Gの色、Bの色、、、という値が沢山集まって一つのモデルになっている、ということですね。

inputができたらそれを使ってonnx形式に変換していきます。これもtorch.onnxモジュールがあるのでそのexportメソッドを叩いてあげればよいだけです。

torch.onnx.export(net, input, output_name, verbose=True, input_names=input_layer_names,
                      output_names=output_layer_names)

torch.onnx.exportメソッドの引数はそれぞれ

  • net: 実行中のモデル
  • input: モデルの入力値
  • output_name: 出力先のパス
  • verbose=True: 変換中の細かいログを吐く
  • input_names: モデルの入力値に関する表示名を指定
  • output_names: モデルの出力値に関する表示名を指定

となります。input_namesやoutput_namesは後から読みやすくするためにつけているだけで、net、input、output_nameがあれば変換は開始されます。

ちなみにこのメソッドではoutput_namesにheatmapとpafという名前を指定しています。ヒートマップというのは検出した人間の関節点を色で表したもので、PAF(Part Affinity Fields)というのは関節同士を連結させた線のようなことを指します。

これでPyTorchのモデルをONNX形式に変換できました。次にこれをOpenVINO形式に変換します。

ONNX => OpenVINO

OpenVINO形式に変換するには、作業環境にOpenVINOのdistributionをインストールする必要があります。他にも色々必要なモジュールがあるので、Dockerで仮想環境を作って、そこで作業したいと思います。

まずこちらよりLinux用のOpenVINO Distributionをダウンロードします。

次にDocker imageを作ってbuildします。こちらがDockerfileになります。

Dockerfile

# Base image with cuda 10.1 with cudnn 7
FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

ENV DEBIAN_FRONTEND noninteractive
ENV LANG C.UTF-8
ENV APT_INSTALL "apt-get install -y --no-install-recommends"

# ==================================================================
# common tools
# ------------------------------------------------------------------

RUN apt-get update && $APT_INSTALL \
    libsm6 \
    libxext6 \
    build-essential \
    apt-utils \
    ca-certificates \
    wget \
    git \
    vim \
    libssl-dev \
    curl \
    unzip \
    unrar \
    python3-dev \
    python3-tk \
    cmake \
    cmake-curses-gui \
    pkg-config \
    sudo \
    wget \
    software-properties-common \
    apt-utils \
    libturbojpeg \
    libxrender1 \
    && rm -rf /var/lib/apt/lists/*



# ==================================================================
# create non-root user
# ------------------------------------------------------------------
ARG USER_ID=1000
RUN useradd -m  -s /bin/bash --no-log-init --system  --uid ${USER_ID} appuser -g sudo
RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
USER appuser
WORKDIR /home/appuser


# ==================================================================
# update pip
# ------------------------------------------------------------------
ENV PATH="/home/appuser/.local/bin:${PATH}"
RUN wget https://bootstrap.pypa.io/get-pip.py && \
	python3 get-pip.py --user && \
	rm get-pip.py


# ==================================================================
# openvino 2020.1
# ------------------------------------------------------------------
ARG DEPENDENCIES="autoconf \
                  automake \
                  build-essential \
                  cmake \
                  cpio \
                  curl \
                  gnupg2 \
                  libdrm2 \
                  libglib2.0-0 \
                  lsb-release \
                  libgtk-3-0 \
                  libtool \
                  udev \
                  unzip"
RUN sudo apt-get update && \
    sudo apt-get install -y --no-install-recommends ${DEPENDENCIES} && \
    sudo rm -rf /var/lib/apt/lists/*

ADD openvino_installers/l_openvino_toolkit_p_2020.4.287.tgz tmp/
RUN sudo chown -R appuser:sudo tmp/
RUN sudo apt-get update && cd tmp/l_openvino_toolkit* && \
    sed -i 's/decline/accept/g' silent.cfg && \
    sudo ./install.sh -s silent.cfg && \
    rm -rf tmp/* && \
    sudo -E /opt/intel/openvino/install_dependencies/install_openvino_dependencies.sh
RUN pip install --user -r /opt/intel/openvino/deployment_tools/model_optimizer/requirements_onnx.txt

GPUを使いたいのでCudaをベースにDocker imageを作ります。
pythonやcmakeなどの基本的なライブラリを入れたら、作業用に[appuser]というユーザを作ります。後はpipとOpenVINOツールキット本体を取り込んだらインストールシェルを叩き、ツールキットの動作に必要なライブラリ群をインストールします。

ADD openvino_installers/l_openvino_toolkit_p_2020.4.287.tgz tmp/

の部分をダウンロードしたツールキットのファイル名に変更しておきましょう。できたらDockerfileと同じ階層に[openvino_installers]というフォルダを作り、その中に先程ダウンロードしたツールキットを入れます。

準備ができたらDockerをbuildします。

docker build -t cuda_10.1_openvino_2020_4:latest .

docker images
docker run --rm -it --ipc=host cuda_10.1_openvino_2020_4 bash

前項でできたonnxファイルをこのコンテナの中にコピーします。

docker ps

docker cp <ダウンロードしたonnxファイルのパス> <docker psで確認したCONTAINER ID>:/tmp

コンテナの中で、先程コピーしたonnxファイルを引数にmodel_optimizerを叩きます。

python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo.py --input_model /tmp/<onnxファイル> --input data --mean_values data[128.0,128.0,128.0] --scale_values data[256] --keep_shape_ops --output stage_1_output_0_pafs,stage_1_output_1_heatmaps

OpenVINOのmodel optimizerはOpenVINO内の推論エンジンで動くような中間表現(Intermediate Representation。IR)に変換する作業を行います。OpenVINOのIRとしては

  • xmlファイル: ネットワークトポロジの定義
  • binファイル: 実データ

の2種類が生成されます。
XMLファイルはモデル構成の把握のために使うので、実際にOpenVINO上で動かすのはbinファイルになりますね。

うまくいけばカレントディレクトリにIRファイルができあがりますのでコンテナからローカルにコピーして終了です。

docker cp <docker psで確認したCONTAINER ID>::/home/appuser/<作成されたbinファイル> ./
docker cp <docker psで確認したCONTAINER ID>::/home/appuser/<作成されたxmlファイル> ./

まとめ

以上、OpenVINO形式への変換方法をまとめてみました。
Chainerなども同じようにonnx形式へ変換して、それ以外の、例えばTensorflowなどはそのままopeimizerにかければOpenVINO形式になります。

OpenVINO形式に変換してOpenVINO上に展開するのもよし、Sagemaker NEOやGCPを使って直接エッジにデプロイするもよし、モデルはいろいろな使い方があります。一つづつ噛み砕いて理解していきましょう。