File size: 3,025 Bytes
9f60e86 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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)
|