🦊

PytorchのモデルをTensorflowjsに変換する

7 min read 2

はじめに

気軽にブラウザで機械学習のモデルを動かしたい、という欲求から、tensorflowjsでいろんなモデルを動かすライブラリを作っています(github)。しかし、下記のグラフに見られるように、ここ最近は、主要カンファレンスの75%のペーパーでPyTorchが使われているようになっており、エッジの効いた面白そうなモデルをtensorflowjsで簡単には動かせない状況になってしまいました。

image
State of AI Report 2020 (site)

ところが、モデル量子化の権威であられるPINTOさんが、今年のAdvent CalenderでPytorchからTensorflow系のモデルへのコンバート手法を公開してくれました(site)。まさに神。ありがたすぎる!

公開された手法を用いてモデルのコンバートを再現してみたところ、私の環境ではいくつか詰まったところがあったので、作業メモとして本記事に残すことにしました。なので、基本はPINTOさんのオリジナルの記事を見ていただければと思います。詰まった場合、もしかしたら参考になるかな、くらいで見ていただければありがたいです。

PyTorch, ONNX, Caffe, OpenVINO (NCHW) のモデルをTensorflow / TensorflowLite (NHWC) へお手軽に変換する

方針

上記で述べたとおり、私の環境が特殊だったためうまくいかないところがありました。自分の環境を壊さない程度にいろいろ試行錯誤したのですがうまく行かなかったので、最終的にはdockerを使って再現成功しました。以下、dockerで作業する前提で話を進めます。

作業内容

docker container起動

今回はCudaのバージョンが10.1のubuntuを使いたいと思います。

$ docker run --gpus all  -v /home/whoami/docker_share:/work --name converter -ti nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

関連パッケージインストール

関連パッケージをガンガンインストールしていきます。
趣味でemacsとかmlocateとか入れてますが、不要なら省いてください。

$ apt update -y
$ apt install -y python3-pip emacs mlocate pciutils cpio git sudo curl
$ python3 -m pip install pip --upgrade
$ pip3 install tensorflow==2.3.1 --upgrade

$ pip3 install openvino2tensorflow --upgrade
$ pip3 install torch==1.7.0+cu101 torchvision==0.8.1+cu101 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
$ pip3 install onnxruntime onnx-simplifier --upgrade
$ pip3 install networkx defusedxml test-generator==0.1.1 tensorflow_datasets tensorflowjs

OpenVino Toolkit インストール

OpenVino Toolkitはpipだけだと入らないようなので、下記のサイトを参考にインストールします。

https://docs.openvinotoolkit.org/latest/openvino_docs_install_guides_installing_openvino_linux.html

このサイトで、OpenVINO toolkit packageを次のサイトからダウンロードするように指示されます。

https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit/download.html

pipでインストールできるようにみえますが、pipでインストールすると、私の環境ではいくつかのコマンドが見つからなくて作業が進められなくなってしまいました。
image

なので、Web and Local Install -> Localを選んでパッケージのダウンロードをしました。registre & Downloadボタンをクリックして、必要な情報を入力してください。
image

Versionは2021.1を選びます。よくわからないので、ちょっと大きいですがFull Packageでダウンロードしておきます。ダウンロードしたら、ファイルをdocker containerから見える場所に置いてください。今回は/work/l_openvino_toolkit_p_2021.1.110.tgzに置きました。

image

それでは、インストールしていきます。
インストールが始まると選択肢が表示されます。いろいろwarningがでますが、Skip prerequisitesでいきます。

$ cp /work/l_openvino_toolkit_p_2021.1.110.tgz ./
$ tar xvfz l_openvino_toolkit_p_2021.1.110.tgz 
$ cd l_openvino_toolkit_p_2021.1.110
$ ./install.sh
$ cd -

次に依存するモジュールをインストールしていきます。

$ cd /opt/intel/openvino_2021/install_dependencies
$ ./install_openvino_dependencies.sh
$ source /opt/intel/openvino_2021/bin/setupvars.sh
$ cd -

以上で完了です!

変換テスト

それでは、早速うまく変換できるか試してみましょう。
PINTOさんの記事を真似てU^2-NetのSemantic Segmentationモデルを変換してみましょう。

U^2-Netのリポジトリをcloneしてモデルをダウンロードしてきます。

$ git clone https://github.com/NathanUA/U-2-Net.git
$ cd U-2-Net/
$ mkdir ./saved_models/u2netp/
$ curl -sc /tmp/cookie "https://drive.google.com/uc?export=download&id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy" > /dev/null
$ CODE="$(awk '/_warning_/ {print $NF}' /tmp/cookie)"  
$ curl -Lb /tmp/cookie "https://drive.google.com/uc?export=download&confirm=${CODE}&id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy" -o saved_models/u2netp/u2netp.pth

以下、変換作業を開始します。

$ export PYTHONPATH=/U-2-Net
$ SIZE=512
$ python3 /opt/intel/openvino_2021/deployment_tools/tools/model_downloader/pytorch_to_onnx.py  \
	--import-module model.u2net   \
	--model-name U2NETP   \
	--input-shape 1,3,${SIZE},${SIZE}   \
	--weights saved_models/u2netp/u2netp.pth   \
	--output-file u2netp_${SIZE}x${SIZE}.onnx   --input-names "x"   \
	--output-names "a/F.sigmoid(d0)"

<snip...>
ONNX check passed successfully.

$ python3 -m onnxsim u2netp_${SIZE}x${SIZE}.onnx u2netp_${SIZE}x${SIZE}_opt.onnx
<snip...>
Checking 2/3...
Ok!

$ python3 /opt/intel/openvino_2021/deployment_tools/model_optimizer/mo.py \
  --input_model u2netp_${SIZE}x${SIZE}_opt.onnx \
  --input_shape [1,3,${SIZE},${SIZE}] \
  --output_dir openvino/${SIZE}x${SIZE}/FP32 \
  --data_type FP32
<snip...>
[ SUCCESS ] Generated IR version 10 model.
[ SUCCESS ] XML file: /U-2-Net/openvino/512x512/FP32/u2netp_512x512.xml
[ SUCCESS ] BIN file: /U-2-Net/openvino/512x512/FP32/u2netp_512x512.bin
[ SUCCESS ] Total execution time: 35.76 seconds. 
[ SUCCESS ] Memory consumed: 809 MB. 


$ openvino2tensorflow \
    --model_path openvino/${SIZE}x${SIZE}/FP32/u2netp_${SIZE}x${SIZE}_opt.xml \
    --model_output_path saved_model_${SIZE}x${SIZE} \
    --output_pb True

<snip...>
All the conversion process is finished! =============================================

$ tensorflowjs_converter \
    --input_format=tf_frozen_model \
    --output_format=tfjs_graph_model \
    --saved_model_tags=serve \
    --output_node_names "Identity" \
    ./saved_model_${SIZE}x${SIZE}/model_float32.pb \
    web_model


これでTensorflowjsのモデルに変換できました。

このモデルを使って実際に画像をセグメンテーションしてみましょう。結果は下記のとおりです。セグメンテーションができていますね。ちょっとボヤけ気味のところもありますが、閾値をうまく設定してあげればもっと良くなるかと思います。
なお、そこそこ軽量という話をよく聞きましたが、私の環境(CPU)では2~3分かかってしまっていました。また、GPUだとRAMが6GBのGTX1660でもメモリ溢れしてしまって動かせませんでした。
(2080Tiも同じマシンに積んでるだけど、ChromeがGTX1660しかActiveにしてくれない。切り替える方法がわからん。。。)
image

次の記事では、U^2-NetのPortrait Drawingについてもやってみます。こんな感じになります。

image

Troubleshoot

ときどき、openvino2tensorflowを実行すると

ModuleNotFoundError: No module named 'openvino.inference_engine'

がでることがある。次のコマンドを実行すると動くようになる。

source /opt/intel/openvino_2021/bin/setupvars.sh

リファレンス

画像は下記のページのものを利用させていただきました。

https://www.pakutaso.com/
https://github.com/NathanUA/U-2-Net

この記事に贈られたバッジ

Discussion

はじめまして。nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04でPython3.6を使ってやってみたのですが、openvino2tensorflowのsaved_modelを生成するところで、エラーになってしまうようです。

質問なのですが、途中でPython3.8をインストールしていますが、使っているのはPython3.6でしょうか?tensorflow 2.3.1 CUDA 10.2をPython3.8でビルドして構築しようとしましたが、openvino-pythonがPython3.8に対応していないように思います。

https://community.intel.com/t5/Intel-Distribution-of-OpenVINO/gError-installing-openvino-2021-1-110/td-p/1217341?profile.language=ja

nabeyangさん
コメントありがとうございます。確かにpython3.8はいらないかもしれないですね。
確認してみたらpipを入れる段階でpython3.6が依存パッケージとして入れられていて、それが使わるようになっていたようです。

ログインするとコメントできます