Merge pull request #711 from marcoschepis/main

Fixed mistake of #710
PINTO0309 · Oct 14, 2024 · 93061bc · 93061bc
2 parents 2ecc03f + 950d780
commit 93061bc
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/README.md b/README.md
@@ -1514,8 +1514,8 @@ usage: onnx2tf
 [-oiqt]
 [-qt {per-channel,per-tensor}]
 [-cind INPUT_NAME NUMPY_FILE_PATH MEAN STD]
-[-iqd {int8,uint8}]
-[-oqd {int8,uint8}]
+[-iqd {int8,uint8,float32}]
+[-oqd {int8,uint8,float32}]
 [-nuo]
 [-nuonag]
 [-b BATCH_SIZE]
@@ -1672,11 +1672,11 @@ optional arguments:
       and {input_op_name}, {numpy_file_path}, {mean}, and {std} must all be entered.
       Otherwise, an error will occur during the -oiqt stage.
 
-  -iqd {int8,uint8}, --input_quant_dtype {int8,uint8}
+  -iqd {int8,uint8,float32}, --input_quant_dtype {int8,uint8,float32}
     Input dtypes when doing Full INT8 Quantization.
     "int8"(default) or "uint8"
 
-  -oqd {int8,uint8}, --output_quant_dtype {int8,uint8}
+  -oqd {int8,uint8,float32}, --output_quant_dtype {int8,uint8,float32}
     Output dtypes when doing Full INT8 Quantization.
     "int8"(default) or "uint8"
 
@@ -2165,11 +2165,11 @@ convert(
 
     input_quant_dtype: Optional[str]
       Input dtypes when doing Full INT8 Quantization.
-      "int8"(default) or "uint8"
+      "int8"(default) or "uint8" or "float32"
 
     output_quant_dtype: Optional[str]
       Output dtypes when doing Full INT8 Quantization.
-      "int8"(default) or "uint8"
+      "int8"(default) or "uint8" or "float32"
 
     not_use_onnxsim: Optional[bool]
       No optimization by onnx-simplifier is performed.

diff --git a/onnx2tf/onnx2tf.py b/onnx2tf/onnx2tf.py
@@ -224,11 +224,11 @@ def convert(
 
     input_quant_dtype: Optional[str]
         Input dtypes when doing Full INT8 Quantization.\n
-        "int8"(default) or "uint8"
+        "int8"(default) or "uint8" or "float32"
 
     output_quant_dtype: Optional[str]
         Output dtypes when doing Full INT8 Quantization.\n
-        "int8"(default) or "uint8"
+        "int8"(default) or "uint8" or "float32"
 
     not_use_onnxsim: Optional[bool]
         No optimization by onnx-simplifier is performed.\n
@@ -1704,13 +1704,17 @@ def representative_dataset_gen():
                     inf_type_input = tf.int8
                 elif input_quant_dtype == 'uint8':
                     inf_type_input = tf.uint8
+                elif input_quant_dtype == 'float32':
+                    inf_type_input = tf.float32
                 else:
                     inf_type_input = tf.int8
-
+                
                 if output_quant_dtype == 'int8':
                     inf_type_output = tf.int8
                 elif output_quant_dtype == 'uint8':
                     inf_type_output = tf.uint8
+                elif output_quant_dtype == 'float32':
+                    inf_type_output = tf.float32
                 else:
                     inf_type_output = tf.int8
                 converter.inference_input_type = inf_type_input
@@ -2144,21 +2148,21 @@ def main():
         '-iqd',
         '--input_quant_dtype',
         type=str,
-        choices=['int8', 'uint8'],
+        choices=['int8', 'uint8', 'float32'],
         default='int8',
         help=\
             'Input dtypes when doing Full INT8 Quantization. \n' +
-            '"int8"(default) or "uint8"'
+            '"int8"(default) or "uint8" or "float32"'
     )
     parser.add_argument(
         '-oqd',
         '--output_quant_dtype',
         type=str,
-        choices=['int8', 'uint8'],
+        choices=['int8', 'uint8', 'float32'],
         default='int8',
         help=\
             'Output dtypes when doing Full INT8 Quantization. \n' +
-            '"int8"(default) or "uint8"'
+            '"int8"(default) or "uint8" or "float32"'
     )
     parser.add_argument(
         '-nuo',