Zenn
🤹

TensorFlow Hubから取得した機械学習モデルをONNX形式に変換する

2023/01/07に公開
Python
TensorFlow
ディープラーニング
機械学習
ONNX
tech

4行まとめ

  • TensorFlow HubからはSavedModel形式で機械学習モデルをダウンロードできる。
  • tf2onnxを使うことでSavedModel形式からONNX形式に変換できる。
  • tf2onnxでは--rename-inputs、--rename-outputsオプションを使うことでノード名を指定できる。
  • onnx、tensorflow、tf2onnxパッケージの依存関係はややこしい。

はじめに

TensorFlow Hubでは様々な機械学習モデルが配布されています。
これらの機械学習モデルをONNX形式に変換して利用する機会がありましたので、その手順を備忘録としてまとめました。

手順

TensorFlow Hubからの機械学習モデルの取得、SavedModel形式からONNX形式への変換の手順は、大きく以下の通りです。

  1. TensorFlow Hubで目的の機械学習モデルを見つける。
  2. SavedModel形式のモデルをダウンロードする。
  3. tar.gzを展開する。
  4. tf2onnxを使ってONNXモデルに変換する。

環境

SavedModel形式からONNX形式への変換にはonnx、tensorflow、tf2onnxパッケージが必要です。
今回は以下のDockerfileを使い、Docker内で作業を実施しました。

FROM ubuntu:22.04
RUN apt-get update \
  && DEBIAN_FRONTEND=noninteractive apt-get install --yes --no-install-recommends \
    ca-certificates \
    python3-dev \
    python3-pip \
    python3-setuptools \
    wget \
  && rm --recursive --force /var/lib/apt/lists/*
RUN python3 -m pip install --upgrade \
  pip==22.3.1 \
  setuptools==65.6.3
RUN python3 -m pip install \
  onnx==1.13.0 \
  tensorflow==2.9.0 \
  tf2onnx==1.13.0

onnx、tensorflow、tf2onnxパッケージの依存関係は幾分ややこしいので、「最新のバージョンを使おう」と考えず、実績のある組み合わせで実施することをお勧めします。

対象の機械学習モデル

今回は以下の機械学習モデルを対象としました。ImageNet-21KでトレーニングされたEfficientNet V2のSサイズ、特徴量抽出用のバージョン2です。

https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet21k_s/feature_vector/2

実際の変換手順

まずはDockerイメージをビルドし、作業ディレクトリをマウントした状態でシェルを起動します。今回はtfhub_onnxというイメージ名にしました。

docker build --tag tfhub_onnx .
docker run --interactive --tty --rm \
  --volume $(pwd):/workspace \
  --workdir /workspace \
  tfhub_onnx

続いて、起動したDockerコンテナ内で変換を実施します。実施例は以下の通りです。

今回使用した機械学習モデルの入力ノード名、出力ノード名は分かりづらいので、 --rename-inputsオプション、--rename-outputsオプションを使って上書きしました。
なお、GPUに関する警告メッセージが出力されていますが、変換するだけならGPUは不要です。

# SavedModel形式の機械学習モデルをダウンロードする

docker$ wget \
  --output-document efficientnet_v2_imagenet21k_s_v2.tar.gz \
  "https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet21k_s/feature_vector/2?tf-hub-format=compressed"

docker$ ls -l efficientnet_v2_imagenet21k_s_v2.tar.gz
-rw-r--r-- 1 root root 76099140 Jul 21  2021 efficientnet_v2_imagenet21k_s_v2.tar.gz

docker$ sha1sum efficientnet_v2_imagenet21k_s_v2.tar.gz
553c076fbc59f9e5a63947297b2d27cbc138efd9  efficientnet_v2_imagenet21k_s_v2.tar.gz

# ダウンロードしたアーカイブを展開する

docker$ mkdir efficientnet_v2_imagenet21k_s_v2
docker$ tar zxfv efficientnet_v2_imagenet21k_s_v2.tar.gz -C efficientnet_v2_imagenet21k_s_v2
./
./variables/
./variables/variables.index
./variables/variables.data-00000-of-00001
./saved_model.pb

# SavedModel形式をONNX形式に変換する

docker$ python3 -m tf2onnx.convert \
  --saved-model efficientnet_v2_imagenet21k_s_v2 \
  --output efficientnet_v2_imagenet21k_s_v2.onnx \
  --rename-inputs input_image \
  --rename-outputs output_feature
2023-01-07 07:13:04.439761: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2023-01-07 07:13:04.439781: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
/usr/lib/python3.10/runpy.py:126: RuntimeWarning: 'tf2onnx.convert' found in sys.modules after import of package 'tf2onnx', but prior to execution of 'tf2onnx.convert'; this may result in unpredictable behaviour
  warn(RuntimeWarning(msg))
2023-01-07 07:13:05.349991: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
2023-01-07 07:13:05.350005: W tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303)
2023-01-07 07:13:05.350029: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:163] no NVIDIA GPU device is present: /dev/nvidia0 does not exist
2023-01-07 07:13:05.350193: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-01-07 07:13:05,351 - WARNING - '--tag' not specified for saved_model. Using --tag serve
2023-01-07 07:13:10,476 - INFO - Signatures found in model: [serving_default].
2023-01-07 07:13:10,476 - WARNING - '--signature_def' not specified, using first signature: serving_default
2023-01-07 07:13:10,478 - INFO - Output names: ['output_1']
2023-01-07 07:13:10.506651: I tensorflow/core/grappler/devices.cc:66] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 0
2023-01-07 07:13:10.506767: I tensorflow/core/grappler/clusters/single_machine.cc:358] Starting new session
2023-01-07 07:13:12.667590: I tensorflow/core/grappler/devices.cc:66] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 0
2023-01-07 07:13:12.667720: I tensorflow/core/grappler/clusters/single_machine.cc:358] Starting new session
2023-01-07 07:13:13,245 - INFO - Using tensorflow=2.9.0, onnx=1.13.0, tf2onnx=1.13.0/2c1db5
2023-01-07 07:13:13,245 - INFO - Using opset <onnx, 13>
2023-01-07 07:13:17,363 - INFO - Computed 0 values for constant folding
2023-01-07 07:13:20,247 - INFO - Optimizing ONNX model
2023-01-07 07:13:24,449 - INFO - After optimization: BatchNormalization -105 (110->5), Const -375 (731->356), GlobalAveragePool +31 (0->31), Identity -2 (2->0), ReduceMean -31 (31->0), Reshape -30 (30->0), Squeeze +1 (0->1), Transpose -589 (590->1)
2023-01-07 07:13:24,513 - INFO -
2023-01-07 07:13:24,513 - INFO - Successfully converted TensorFlow model efficientnet_v2_imagenet21k_s_v2 to ONNX
2023-01-07 07:13:24,513 - INFO - Model inputs: ['input_image']
2023-01-07 07:13:24,513 - INFO - Model outputs: ['output_feature']
2023-01-07 07:13:24,513 - INFO - ONNX model is saved at efficientnet_v2_imagenet21k_s_v2.onnx

docker$ ls -l efficientnet_v2_imagenet21k_s_v2.onnx
-rw-r--r-- 1 root root 80650295 Jan  7 07:13 efficientnet_v2_imagenet21k_s_v2.onnx

おわりに

パッケージの依存関係でハマり時間を浪費してしまいましたが、無事にSavedModel形式をONNX形式に変換することができました。
何らかの参考になれば幸いです。

Discussion