🦊

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

2020/12/13に公開

2件

JavaScript

TensorFlow

tech

はじめに

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

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

OpenVino Toolkit インストール

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

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

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

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

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

それでは、インストールしていきます。
インストールが始まると選択肢が表示されます。いろいろ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にしてくれない。切り替える方法がわからん。。。)

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

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

nabeyang

はじめまして。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に対応していないように思います。

wok

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