tftrt convert option changed

PINTO0309 · Oct 10, 2021 · 98a146b · 98a146b
1 parent dcb043b
commit 98a146b
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 19 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -7,7 +7,7 @@ ARG CPVER=cp38
 ARG OPENVINOVER=2021.4.582
 ARG OPENVINOROOTDIR=/opt/intel/openvino_2021
 ARG TENSORRTVER=cuda11.3-trt8.0.1.6-ga-20210626
-ARG APPVER=v1.12.1
+ARG APPVER=v1.13.0
 ARG wkdir=/home/user
 
 # dash -> bash

diff --git a/README.md b/README.md
@@ -296,7 +296,8 @@ usage: tflite2tensorflow
   [--tfds_download_flg]
   [--load_dest_file_path_for_the_calib_npy LOAD_DEST_FILE_PATH_FOR_THE_CALIB_NPY]
   [--output_tfjs]
-  [--output_tftrt]
+  [--output_tftrt_float32]
+  [--output_tftrt_float16]
   [--output_coreml]
   [--output_edgetpu]
   [--edgetpu_compiler_timeout EDGETPU_COMPILER_TIMEOUT]
@@ -366,8 +367,10 @@ optional arguments:
                         [20, 513, 513, 3] -> [Number of images, h, w, c]
   --output_tfjs
                         tfjs model output switch
-  --output_tftrt
-                        tftrt model output switch
+  --output_tftrt32
+                        tftrt float32 model output switch
+  --output_tftrt16
+                        tftrt float16 model output switch
   --output_coreml
                         coreml model output switch
   --output_edgetpu
@@ -440,7 +443,8 @@ $ tflite2tensorflow \
   --string_formulas_for_normalization 'data / 255.0' \
   --output_tfjs \
   --output_coreml \
-  --output_tftrt \
+  --output_tftrt_float32 \
+  --output_tftrt_float16 \
   --output_onnx \
   --onnx_opset 13 \
   --output_openvino_and_myriad
@@ -460,7 +464,8 @@ $ tflite2tensorflow \
   --string_formulas_for_normalization 'data / 255.0' \
   --output_tfjs \
   --output_coreml \
-  --output_tftrt \
+  --output_tftrt_float32 \
+  --output_tftrt_float16 \
   --output_onnx \
   --onnx_opset 13
 ```

diff --git a/setup.py b/setup.py
@@ -11,7 +11,7 @@
     setup(
         name="tflite2tensorflow",
         scripts=scripts,
-        version="1.12.1",
+        version="1.13.0",
         description="Generate saved_model, tfjs, tf-trt, EdgeTPU, CoreML, quantized tflite, ONNX, OpenVINO, Myriad Inference Engine blob and .pb from .tflite.",
         long_description=long_description,
         long_description_content_type="text/markdown",

diff --git a/tflite2tensorflow/tflite2tensorflow.py b/tflite2tensorflow/tflite2tensorflow.py
@@ -5372,7 +5372,8 @@ def main():
     npy_load_default_path = 'sample_npy/calibration_data_img_sample.npy'
     parser.add_argument('--load_dest_file_path_for_the_calib_npy', type=str, default=npy_load_default_path, help='The path from which to load the .npy file containing the numpy binary version of the calibration data. Default: sample_npy/calibration_data_img_sample.npy')
     parser.add_argument('--output_tfjs', action='store_true', help='tfjs model output switch')
-    parser.add_argument('--output_tftrt', action='store_true', help='tftrt model output switch')
+    parser.add_argument('--output_tftrt_float32', action='store_true', help='tftrt float32 model output switch')
+    parser.add_argument('--output_tftrt_float16', action='store_true', help='tftrt float16 model output switch')
     parser.add_argument('--output_coreml', action='store_true', help='coreml model output switch')
     parser.add_argument('--output_edgetpu', action='store_true', help='edgetpu model output switch')
     parser.add_argument('--edgetpu_compiler_timeout', type=int, default=3600, help='edgetpu_compiler timeout for one compilation process in seconds. Default: 3600')
@@ -5414,7 +5415,8 @@ def main():
     tfds_download_flg = args.tfds_download_flg
     load_dest_file_path_for_the_calib_npy = args.load_dest_file_path_for_the_calib_npy
     output_tfjs = args.output_tfjs
-    output_tftrt = args.output_tftrt
+    output_tftrt_float32 = args.output_tftrt_float32
+    output_tftrt_float16 = args.output_tftrt_float16
     output_coreml = args.output_coreml
     output_edgetpu = args.output_edgetpu
     edgetpu_compiler_timeout = args.edgetpu_compiler_timeout
@@ -5452,7 +5454,7 @@ def main():
             print('\'tensorflowjs\' is not installed. Please run the following command to install \'tensorflowjs\'.')
             print('pip3 install --upgrade tensorflowjs')
             sys.exit(-1)
-    if output_tftrt:
+    if output_tftrt_float32 or output_tftrt_float16:
         if not 'tensorrt' in package_list:
             print('\'tensorrt\' is not installed. Please check the following website and install \'tensorrt\'.')
             print('https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html')
@@ -5510,17 +5512,18 @@ def main():
                     output_integer_quant_tflite or \
                         output_full_integer_quant_tflite or \
                             output_tfjs or \
-                                output_tftrt or \
-                                    output_coreml or \
-                                        output_edgetpu or \
-                                            output_onnx or \
-                                                output_openvino_and_myriad:
+                                output_tftrt_float32 or \
+                                    output_tftrt_float16 or \
+                                        output_coreml or \
+                                            output_edgetpu or \
+                                                output_onnx or \
+                                                    output_openvino_and_myriad:
         tfv2_flg = True
 
     if tfv1_flg and tfv2_flg:
         print(f'{Color.RED}ERROR:{Color.RESET} Group.1 and Group.2 cannot be set to True at the same time. Please specify either Group.1 or Group.2.')
         print('[Group.1] output_pb')
-        print('[Group.2] output_no_quant_float32_tflite, output_weight_quant_tflite, output_float16_quant_tflite, output_integer_quant_tflite, output_full_integer_quant_tflite, output_tfjs, output_tftrt, output_coreml, output_edgetpu, output_onnx, output_openvino_and_myriad')
+        print('[Group.2] output_no_quant_float32_tflite, output_weight_quant_tflite, output_float16_quant_tflite, output_integer_quant_tflite, output_full_integer_quant_tflite, output_tfjs, output_tftrt_float32, output_tftrt_float16, output_coreml, output_edgetpu, output_onnx, output_openvino_and_myriad')
         sys.exit(-1)
 
     if optimizing_for_openvino_and_myriad and optimizing_hardswish_for_edgetpu:
@@ -5919,7 +5922,7 @@ def representative_dataset_gen():
                 traceback.print_exc()
 
         # TF-TRT (TensorRT) convert
-        if output_tftrt:
+        if output_tftrt_float32:
             try:
                 def input_fn():
                     input_shapes = []
@@ -5943,8 +5946,22 @@ def input_fn():
                         if input_shapes_permutations_idx > len(input_shapes_permutations):
                             print(f'{Color.RED}ERROR:{Color.RESET}', e)
 
-                converter.save(f'{model_output_path}/tensorrt_saved_model_float32')
-                print(f'{Color.GREEN}TF-TRT (TensorRT) convertion complete!{Color.RESET} - {model_output_path}/tensorrt_saved_model_float32')
+            except Exception as e:
+                print(f'{Color.RED}ERROR:{Color.RESET}', e)
+                import traceback
+                traceback.print_exc()
+                print(f'{Color.RED}The binary versions of TensorFlow and TensorRT may not be compatible. Please check the version compatibility of each package.{Color.RESET}')
+
+        # TF-TRT (TensorRT) convert
+        if output_tftrt_float16:
+            try:
+                def input_fn():
+                    input_shapes = []
+                    for shape in input_shapes_permutations[input_shapes_permutations_idx]:
+                        shape_tuple = tuple(shape)
+                        input_shapes.append(np.zeros(shape_tuple).astype(np.float32))
+                    yield input_shapes
+
                 print(f'{Color.REVERCE}TF-TRT (TensorRT) Float16 convertion started{Color.RESET}', '=' * 40)
                 params = tf.experimental.tensorrt.ConversionParams(precision_mode='FP16', maximum_cached_engines=10000)
                 converter = tf.experimental.tensorrt.Converter(input_saved_model_dir=model_output_path, conversion_params=params)