Add utility to convert ONNX model in FP32/16 mixed precision

convert_to_mixed.py

@@ -0,0 +1,68 @@
+import sys
+import numpy as np
+from PIL import Image
+import onnx
+import onnxruntime as ort
+from onnxconverter_common import auto_mixed_precision_model_path
+import argparse
+
+PROVIDERS=[('TensorrtExecutionProvider', {'trt_fp16_enable':True,}), 'CUDAExecutionProvider', 'CPUExecutionProvider']
+RTOL=0.1
+ATOL=0.1
+
+def detect_model_input_size(model_path):
+    model = onnx.load(model_path)
+    for input_tensor in model.graph.input:
+        # Assuming the input node is named 'input'
+        if input_tensor.name == 'input':
+            tensor_shape = input_tensor.type.tensor_type.shape
+            # Extract the dimensions: (batch_size, channels, height, width)
+            dims = [dim.dim_value for dim in tensor_shape.dim]
+            # Replace dynamic batch size (-1 or 0) with 1
+            if dims[0] < 1:
+                dims[0] = 1
+            return tuple(dims[2:4])  # Return (height, width)
+    raise ValueError("Input node 'input' not found in the model")
+
+def load_and_preprocess_image(image_path, size=(224, 224)):
+    image = Image.open(image_path).convert('RGB')
+    image = image.resize(size)
+    image = np.array(image).astype(np.float32) / 255.
+    image = np.transpose(image, (2, 0, 1))
+    image = np.expand_dims(image, axis=0)
+    return image
+
+def infer(model_path, input_feed):
+    session = ort.InferenceSession(model_path, providers=PROVIDERS)
+    input_name = session.get_inputs()[0].name
+    result = session.run(None, {input_name: input_feed})
+    return result
+
+def main(args):
+    model_input_size = detect_model_input_size(args.source_model_path)
+    input_feed = {'input':load_and_preprocess_image(args.test_image_path, size=model_input_size)}
+
+    auto_mixed_precision_model_path.auto_convert_mixed_precision_model_path(source_model_path=args.source_model_path, 
+                                    input_feed=input_feed, 
+                                    target_model_path=args.target_model_path,
+                                    customized_validate_func=None, 
+                                    rtol=RTOL, atol=ATOL,
+                                    provider=PROVIDERS, 
+                                    keep_io_types=True,
+                                    verbose=True)
+
+    original_result = infer(args.source_model_path, input_feed)
+    converted_result = infer(args.target_model_path, input_feed)
+    
+    is_close = np.allclose(original_result[0], converted_result[0], rtol=RTOL, atol=ATOL)
+    print(f"Validation result: {'Success' if is_close else 'Failure'}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert an ONNX model to mixed precision format.")
+    parser.add_argument("source_model_path", type=str, help="Path to the source ONNX model.")
+    parser.add_argument("target_model_path", type=str, help="Path where the mixed precision model will be saved.")
+    parser.add_argument("test_image_path", type=str, help="Path to a test image for validating the model conversion.")
+
+    args = parser.parse_args()
+    
+    main(args)