Compatible with ONNX model (NHWC) output

PINTO0309 · Feb 8, 2021 · 134665f · 134665f
1 parent e028b4d
commit 134665f
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -29,27 +29,31 @@ Work in progress now.
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TF-TRT (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> EdgeTPU (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> CoreML (NHWC)
+ - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> ONNX (NHWC)
 
 - Caffe (NCHW) -> OpenVINO (NCHW) ->
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TFLite (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TFJS (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TF-TRT (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> EdgeTPU (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> CoreML (NHWC)
+ - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> ONNX (NHWC)
 
 - MXNet (NCHW) -> OpenVINO (NCHW) ->
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TFLite (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TFJS (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TF-TRT (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> EdgeTPU (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> CoreML (NHWC)
+ - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> ONNX (NHWC)
 
 - Keras (NHWC) -> OpenVINO (NCHW・Optimized) ->
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TFLite (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TFJS (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> TF-TRT (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> EdgeTPU (NHWC)
  - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> CoreML (NHWC)
+ - -> **`openvino2tensorflow`** -> Tensorflow/Keras (NHWC) -> ONNX (NHWC)
 
 - saved_model -> **`saved_model_to_pb`** -> pb
 
@@ -150,7 +154,7 @@ Work in progress now.
 |84|DepthToSpace|DepthToSpace||
 |85|Sqrt|sqrt||
 |86|SquaredDifference|squared_difference||
-|87|FakeQuantize|subtract, multiply, round, greater, where, less_equal, add|experimental|
+|87|FakeQuantize|subtract, multiply, round, greater, where, less_equal, add||
 |88|Result|Identity|Output|
 
 ## 4. Setup
@@ -192,6 +196,7 @@ usage: openvino2tensorflow [-h] --model_path MODEL_PATH
  [--output_tftrt OUTPUT_TFTRT]
  [--output_coreml OUTPUT_COREML]
  [--output_edgetpu OUTPUT_EDGETPU]
+ [--output_onnx OUTPUT_ONNX]
  [--replace_swish_and_hardswish REPLACE_SWISH_AND_HARDSWISH]
  [--optimizing_hardswish_for_edgetpu OPTIMIZING_HARDSWISH_FOR_EDGETPU]
  [--replace_prelu_and_minmax REPLACE_PRELU_AND_MINMAX]
@@ -255,6 +260,8 @@ optional arguments:
  coreml model output switch
  --output_edgetpu OUTPUT_EDGETPU
  edgetpu model output switch
+ --output_onnx OUTPUT_ONNX
+ onnx model output switch
  --replace_swish_and_hardswish REPLACE_SWISH_AND_HARDSWISH
  Replace swish and hard-swish with each other
  --optimizing_hardswish_for_edgetpu OPTIMIZING_HARDSWISH_FOR_EDGETPU
@@ -294,6 +301,7 @@ usage: saved_model_to_tflite [-h] --saved_model_dir_path
  [--output_tftrt OUTPUT_TFTRT]
  [--output_coreml OUTPUT_COREML]
  [--output_edgetpu OUTPUT_EDGETPU]
+ [--output_onnx OUTPUT_ONNX]
 
 optional arguments:
  -h, --help show this help message and exit
@@ -349,6 +357,8 @@ optional arguments:
  coreml model output switch
  --output_edgetpu OUTPUT_EDGETPU
  edgetpu model output switch
+ --output_onnx OUTPUT_ONNX
+ onnx model output switch
 ```
 ### 5-3. pb to saved_model convert
 ```bash

diff --git a/openvino2tensorflow/openvino2tensorflow.py b/openvino2tensorflow/openvino2tensorflow.py
@@ -113,6 +113,7 @@ def convert(model,
  output_tftrt,
  output_coreml,
  output_edgetpu,
+ output_onnx,
  replace_swish_and_hardswish,
  optimizing_hardswish_for_edgetpu,
  replace_prelu_and_minmax,
@@ -342,7 +343,6 @@ def parse_json(jsonfile_path):
  del added_key_list
 
  layer_id_port_dict = get_num_of_outputs_per_layer_id(tf_edges)
- # print(layer_id_port_dict)
 
  # for i in tf_edges.items():
  # print(i)
@@ -396,15 +396,6 @@ def parse_json(jsonfile_path):
  else:
  tf_layers_dict[layer_id] = decodedwgt
 
- # #############################################################
- # if layer_id == '1123':
- # tf_layers_dict[layer_id] = np.array([1,2,513,513])
- # elif layer_id == '1125':
- # tf_layers_dict[layer_id] = np.array([0,3,1,2])
- # else:
- # tf_layers_dict[layer_id] = decodedwgt
- # #############################################################
-
  ### Convolution
  elif layer.attrib['type'] == 'Convolution':
  # port0 = [int(sdim.text) for sdim in layer.find('input')[0]]
@@ -2202,6 +2193,23 @@ def input_fn():
  print('Please install edgetpu_compiler according to the following website.')
  print('https://coral.ai/docs/edgetpu/compiler/#system-requirements')
 
+ # ONNX convert
+ if output_onnx:
+ import subprocess
+ try:
+ print(f'{Color.REVERCE}ONNX convertion started{Color.RESET}', '=' * 61)
+ result = subprocess.check_output(['python3',
+ '-m', 'tf2onnx.convert',
+ '--saved-model', model_output_path,
+ '--output', f'{model_output_path}/model_float32.onnx'],
+ stderr=subprocess.PIPE).decode('utf-8')
+ print(result)
+ print(f'{Color.GREEN}ONNX convertion complete!{Color.RESET} - {model_output_path}/tfjs_model_float32')
+ except subprocess.CalledProcessError as e:
+ print(f'{Color.RED}ERROR:{Color.RESET}', e.stderr.decode('utf-8'))
+ import traceback
+ traceback.print_exc()
+
 def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--model_path', type=str, required=True, help='input IR model path (.xml)')
@@ -2227,6 +2235,7 @@ def main():
  parser.add_argument('--output_tftrt', type=bool, default=False, help='tftrt model output switch')
  parser.add_argument('--output_coreml', type=bool, default=False, help='coreml model output switch')
  parser.add_argument('--output_edgetpu', type=bool, default=False, help='edgetpu model output switch')
+ parser.add_argument('--output_onnx', type=bool, default=False, help='onnx model output switch')
  parser.add_argument('--replace_swish_and_hardswish', type=bool, default=False, help='Replace swish and hard-swish with each other')
  parser.add_argument('--optimizing_hardswish_for_edgetpu', type=bool, default=False, help='Optimizing hardswish for edgetpu')
  parser.add_argument('--replace_prelu_and_minmax', type=bool, default=False, help='Replace prelu and minimum/maximum with each other')
@@ -2260,6 +2269,7 @@ def main():
  output_tftrt = args.output_tftrt
  output_coreml = args.output_coreml
  output_edgetpu = args.output_edgetpu
+ output_onnx = args.output_onnx
  replace_swish_and_hardswish = args.replace_swish_and_hardswish
  optimizing_hardswish_for_edgetpu = args.optimizing_hardswish_for_edgetpu
  replace_prelu_and_minmax = args.replace_prelu_and_minmax
@@ -2279,7 +2289,8 @@ def main():
  not output_tfjs and \
  not output_tftrt and \
  not output_coreml and \
- not output_edgetpu:
+ not output_edgetpu and \
+ not output_onnx:
  print('Set at least one of the output switches (output_*) to true.')
  sys.exit(-1)
 
@@ -2306,6 +2317,12 @@ def main():
  print('\'coremltoos\' is not installed. Please run the following command to install \'coremltoos\'.')
  print('pip3 install --upgrade coremltools')
  sys.exit(-1)
+ if output_onnx:
+ if not 'tf2onnx' in package_list:
+ print('\'tf2onnx\' is not installed. Please run the following command to install \'tf2onnx\'.')
+ print('pip3 install --upgrade tf2onnx')
+ sys.exit(-1)
+
  if output_integer_quant_tflite or output_full_integer_quant_tflite:
  if not 'tensorflow-datasets' in package_list:
  print('\'tensorflow-datasets\' is not installed. Please run the following command to install \'tensorflow-datasets\'.')
@@ -2339,7 +2356,7 @@ def main():
  string_formulas_for_normalization,
  calib_ds_type, ds_name_for_tfds_for_calibration, split_name_for_tfds_for_calibration,
  download_dest_folder_path_for_the_calib_tfds, tfds_download_flg,
- output_tfjs, output_tftrt, output_coreml, output_edgetpu,
+ output_tfjs, output_tftrt, output_coreml, output_edgetpu, output_onnx,
  replace_swish_and_hardswish, optimizing_hardswish_for_edgetpu, replace_prelu_and_minmax,
  yolact, weight_replacement_config, debug, debug_layer_number)
  print(f'{Color.REVERCE}All the conversion process is finished!{Color.RESET}', '=' * 45)

diff --git a/openvino2tensorflow/saved_model_to_tflite.py b/openvino2tensorflow/saved_model_to_tflite.py
@@ -52,7 +52,8 @@ def convert(saved_model_dir_path,
  output_tfjs,
  output_tftrt,
  output_coreml,
- output_edgetpu):
+ output_edgetpu,
+ output_onnx):
 
  print(f'{Color.REVERCE}Start conversion process from saved_model to tflite{Color.RESET}', '=' * 38)
 
@@ -293,6 +294,23 @@ def input_fn():
  import traceback
  traceback.print_exc()
 
+ # ONNX convert
+ if output_onnx:
+ import subprocess
+ try:
+ print(f'{Color.REVERCE}ONNX convertion started{Color.RESET}', '=' * 61)
+ result = subprocess.check_output(['python3',
+ '-m', 'tf2onnx.convert',
+ '--saved-model', model_output_dir_path,
+ '--output', f'{model_output_dir_path}/model_float32.onnx'],
+ stderr=subprocess.PIPE).decode('utf-8')
+ print(result)
+ print(f'{Color.GREEN}ONNX convertion complete!{Color.RESET} - {model_output_dir_path}/tfjs_model_float32')
+ except subprocess.CalledProcessError as e:
+ print(f'{Color.RED}ERROR:{Color.RESET}', e.stderr.decode('utf-8'))
+ import traceback
+ traceback.print_exc()
+
 def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--saved_model_dir_path', type=str, required=True, help='Input saved_model dir path')
@@ -316,6 +334,7 @@ def main():
  parser.add_argument('--output_tftrt', type=bool, default=False, help='tftrt model output switch')
  parser.add_argument('--output_coreml', type=bool, default=False, help='coreml model output switch')
  parser.add_argument('--output_edgetpu', type=bool, default=False, help='edgetpu model output switch')
+ parser.add_argument('--output_onnx', type=bool, default=False, help='onnx model output switch')
 
  args = parser.parse_args()
  saved_model_dir_path = args.saved_model_dir_path
@@ -348,6 +367,7 @@ def main():
  output_tftrt = args.output_tftrt
  output_coreml = args.output_coreml
  output_edgetpu = args.output_edgetpu
+ output_onnx = args.output_onnx
 
  if not output_no_quant_float32_tflite and \
  not output_weight_quant_tflite and \
@@ -356,7 +376,8 @@ def main():
  not output_tfjs and \
  not output_tftrt and \
  not output_coreml and \
- not output_edgetpu:
+ not output_edgetpu and \
+ not output_onnx:
  print('Set at least one of the output switches (output_*) to true.')
  sys.exit(-1)
 
@@ -383,6 +404,11 @@ def main():
  print('\'coremltoos\' is not installed. Please run the following command to install \'coremltoos\'.')
  print('pip3 install --upgrade coremltools')
  sys.exit(-1)
+ if output_onnx:
+ if not 'tf2onnx' in package_list:
+ print('\'tf2onnx\' is not installed. Please run the following command to install \'tf2onnx\'.')
+ print('pip3 install --upgrade tf2onnx')
+ sys.exit(-1)
 
  if output_integer_quant_tflite or output_full_integer_quant_tflite:
  if not 'tensorflow-datasets' in package_list:
@@ -426,7 +452,8 @@ def main():
  output_tfjs,
  output_tftrt,
  output_coreml,
- output_edgetpu)
+ output_edgetpu,
+ output_onnx)
  print(f'{Color.REVERCE}All the conversion process is finished!{Color.RESET}', '=' * 45)
 
 if __name__ == "__main__":

diff --git a/setup.py b/setup.py
@@ -11,7 +11,7 @@
  setup(
  name="openvino2tensorflow",
  scripts=scripts,
- version="1.7.0",
+ version="1.8.0",
  description="This script converts the OpenVINO IR model to Tensorflow's saved_model, tflite, h5 and pb. in (NCHW) format",
  long_description=long_description,
  long_description_content_type="text/markdown",