🤪

TensorFlow 2 Detection Model Zoo モデルのtflite変換

に公開3

1. はじめに

TensorFlow 2 Detection Model Zoo というリポジトリがあります。 こちらには TensorFlow v2 ベースのトレーニング済みの様々な物体検出モデルがコミットされているのですが、正規の手順、というか、確立された手順がどこにも言及されていないため、自力で適当に変換しました。手順をシェアします。秒殺で終わるので完全に書き捨ての記事です。
https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md

2. 手順

トレーニング済みのモデルをダウンロードして解凍します。

$ wget http://download.tensorflow.org/models/object_detection/tf2/20200711/centernet_resnet50_v2_512x512_kpts_coco17_tpu-8.tar.gz
$ tar -zxvf centernet_resnet50_v2_512x512_kpts_coco17_tpu-8.tar.gz
$ rm centernet_resnet50_v2_512x512_kpts_coco17_tpu-8.tar.gz
$ cd centernet_resnet50_v2_512x512_kpts_coco17_tpu-8

ダウンロードして解凍したときに同梱されている saved_model の入力と出力の名前や形状を確認します。下記のコマンドを実行します。

$ saved_model_cli show \
  --dir saved_model \
  --tag_set serve \
  --signature_def serving_default

The given SavedModel SignatureDef contains the following input(s):
  inputs['input_tensor'] tensor_info:
      dtype: DT_UINT8
      shape: (1, -1, -1, 3)
      name: serving_default_input_tensor:0
The given SavedModel SignatureDef contains the following output(s):
  outputs['detection_boxes'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100, 4)
      name: StatefulPartitionedCall:0
  outputs['detection_classes'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100)
      name: StatefulPartitionedCall:1
  outputs['detection_keypoint_scores'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100, 17)
      name: StatefulPartitionedCall:2
  outputs['detection_keypoints'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100, 17, 2)
      name: StatefulPartitionedCall:3
  outputs['detection_scores'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100)
      name: StatefulPartitionedCall:4
  outputs['num_detections'] tensor_info:
      dtype: DT_FLOAT
      shape: (1)
      name: StatefulPartitionedCall:5
Method name is: tensorflow/serving/predict

下記のロジックを作成して test.py として保存します。 concrete function というものを利用して saved_model の入力解像度 [1, -1, -1, 3][1, 320, 320, 3] に固定します。変更先の解像度は自由に変更してください。tfliteはバッチサイズ以外の次元が不定の状態 (-1) のままだと変換に失敗します。 なお、入力のTensorが複数ある場合は、 input_shapes = [[1,320,320,3]] の部分を input_shapes = [[1,320,320,3],[1,10,10,1]] のようにリストでつなげて記載するだけです。

### tf-nightly==2.6.0-dev20210430

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph

# Load saved_model and change input shape
# https://github.com/tensorflow/tensorflow/issues/30180#issuecomment-505959220
model = tf.saved_model.load('saved_model')
input_shapes = [[1,320,320,3]]

concrete_func = \
    model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
concrete_func_input_tensors = \
    [tensor for tensor in concrete_func.inputs if tensor.dtype != tf.resource and not 'unknown' in tensor.name]

for conc_input, def_input in zip(concrete_func_input_tensors, input_shapes):
    print('Before changing the input shape', conc_input)
    conc_input.set_shape(def_input)
    print('After  changing the input shape', conc_input)

converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter.target_ops = \
    [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()
open("test.tflite", "wb").write(tflite_model)

TensorFlowの環境をクリーニングします。ご自身の手元の環境を壊したくない場合は、こちらのDockerコンテナをご利用ください。 Docker HubからPullして何度でも破壊できる環境がすぐに利用できます。
https://github.com/PINTO0309/tflite2tensorflow

$ sudo pip3 uninstall -y tensorboard-plugin-wit tb-nightly tensorboard \
                      tf-estimator-nightly tensorflow-gpu \
                      tensorflow tf-nightly tensorflow_estimator

tf-nightly を導入します。

$ sudo pip3 install tf-nightly==2.6.0-dev20210430

変換スクリプトを実行します。

$ python3 test.py

終わり。簡単ですね。 test.tflite というファイルが生成されているのが確認できます。
result

Discussion

Shinobu HUYUGIRIShinobu HUYUGIRI

はじめまして!

BigGAN(TensorFlow 1)をbatch_sizeが固定されたtfliteへ変換したく、情報を探していたところ、この記事に辿り着きました。私の手元では、TensorFlow 1のせいなのか tflite_model = converter.convert() の実行時に現状以下のエラーが発生してしまっているため、以降はTensorFlowのコミュニティで質問してみようと思います。

変換するための手がかりを掴めて嬉しいです。
有用な情報を本当にありがとうございます!

▼変換したいモデル
https://www.kaggle.com/models/deepmind/biggan

Before changing the input shape Tensor("truncation:0", shape=(), dtype=float32)
After  changing the input shape Tensor("truncation:0", shape=(), dtype=float32)
Before changing the input shape Tensor("y:0", shape=(None, 1000), dtype=float32)
After  changing the input shape Tensor("y:0", shape=(1, 1000), dtype=float32)
Before changing the input shape Tensor("z:0", shape=(None, 120), dtype=float32)
After  changing the input shape Tensor("z:0", shape=(1, 120), dtype=float32)
...
Input 0 of node cond/AssignVariableOp was passed float from cond/AssignVariableOp/Switch:0 incompatible with expected resource.
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input 0 of node cond/AssignVariableOp was passed float from cond/AssignVariableOp/Switch:0 incompatible with expected resource.

During handling of the above exception, another exception occurred:

  File "/home/shino/sandbox/python/biggan/biggan_tflite.py", line 21, in <module>
    tflite_model = converter.convert()
ValueError: Input 0 of node cond/AssignVariableOp was passed float from cond/AssignVariableOp/Switch:0 incompatible with expected resource.
PINTOPINTO

変換後の精度は保証しませんが、変換自体は簡単にできます。頑張ってください。

pip install -U tf2onnx

python -m tf2onnx.convert \
--saved-model . \
--tag '' \
--output biggan-tensorflow1-128-v2.onnx

pip install onnxsim==0.4.30

ONNXSIM_FIXED_POINT_ITERS=10000 onnxsim biggan-tensorflow1-128-v2.onnx biggan-tensorflow1-128-v2.onnx \
--overwrite-input-shape "y:1,1000" "z:1,120"

pip install -U onnx==1.17.0 \
&& pip install -U nvidia-pyindex \
&& pip install -U onnx-graphsurgeon \
&& pip install -U onnxruntime==1.18.1 \
&& pip install -U onnxsim==0.4.33 \
&& pip install -U simple_onnx_processing_tools \
&& pip install -U sne4onnx>=1.0.13 \
&& pip install -U sng4onnx>=1.0.4 \
&& pip install -U ai_edge_litert==1.2.0 \
&& pip install -U tensorflow==2.19.0 \
&& pip install -U protobuf==3.20.3 \
&& pip install -U onnx2tf \
&& pip install -U h5py==3.11.0 \
&& pip install -U psutil==5.9.5 \
&& pip install -U ml_dtypes==0.5.1 \
&& pip install -U tf-keras==2.19.0 \
&& pip install flatbuffers>=23.5.26
replace.json
{
  "format_version": 1,
  "operations": [
    {
      "op_name": "Generator_2/GBlock/DepthToSpace__35",
      "param_target": "attributes",
      "param_name": "perm",
      "values": [0,1,2,3]
    },
    {
      "op_name": "Generator_2/GBlock/DepthToSpace_1__27",
      "param_target": "attributes",
      "param_name": "perm",
      "values": [0,1,2,3]
    },
    {
      "op_name": "Generator_2/GBlock_1/DepthToSpace__55",
      "param_target": "attributes",
      "param_name": "perm",
      "values": [0,1,2,3]
    },
    {
      "op_name": "Generator_2/GBlock_2/DepthToSpace__75",
      "param_target": "attributes",
      "param_name": "perm",
      "values": [0,1,2,3]
    },
    {
      "op_name": "Generator_2/GBlock_3/DepthToSpace__95",
      "param_target": "attributes",
      "param_name": "perm",
      "values": [0,1,2,3]
    },
    {
      "op_name": "Generator_2/GBlock_4/DepthToSpace__145",
      "param_target": "attributes",
      "param_name": "perm",
      "values": [0,1,2,3]
    }
  ]
}
onnx2tf -i biggan-tensorflow1-128-v2.onnx -cotof -prf replace.json

Shinobu HUYUGIRIShinobu HUYUGIRI

お忙しいところ
本当にありがとうございます!!

onnx2tfのJSONの書き方に
悩んでいたため、とても助かりました。

頑張ってみます!!