convert_to_mixed.py

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)