From a201a55da672295aa22adc69a69b57dd34b94643 Mon Sep 17 00:00:00 2001 From: shai Date: Tue, 19 Mar 2024 18:02:38 -0500 Subject: [PATCH] update onnx export script --- .../detection/_utils/export_det_onnx.py | 125 +++++++++++------- 1 file changed, 79 insertions(+), 46 deletions(-) diff --git a/doctr/models/detection/_utils/export_det_onnx.py b/doctr/models/detection/_utils/export_det_onnx.py index ee5cd6874c..4d634af0fc 100644 --- a/doctr/models/detection/_utils/export_det_onnx.py +++ b/doctr/models/detection/_utils/export_det_onnx.py @@ -1,62 +1,95 @@ +# Import required libraries import time - +import torch import numpy as np -import torch.onnx - +import onnxruntime from doctr.models import ocr_predictor +from openvino.runtime import Core -model = ocr_predictor(det_arch ='db_resnet50_rotation', pretrained=True) -model.det_predictor.model = model.det_predictor.model.eval() -input = torch.randn(1, 3, 1024, 1024) -input2 = torch.randn(1, 3, 1536, 1536) -start = time.time() -pred = model.det_predictor.model(input) -print("pytorch time", time.time() - start) -torch.onnx.export(model.det_predictor.model, - input, - 'db_resnet50_rotation.onnx', - export_params = True, +start_load_time = time.time() +device = torch.device('cpu') +model = ocr_predictor(det_arch='db_resnet50', pretrained=True).det_predictor.model +model.to(device).eval() +model_load_time = time.time() - start_load_time +print(f"PyTorch Model Load Time: {model_load_time} seconds") + + +# Define a function for PyTorch inference and benchmarking +def pytorch_inference(model, input_tensor): + with torch.no_grad(): + return model(input_tensor).detach().cpu().numpy() + + +# Define a function to benchmark ONNX inference and verify accuracy +def benchmark_onnx_inference_and_verify(model_path, input_tensor, pytorch_output): + session = onnxruntime.InferenceSession(model_path, providers=["CPUExecutionProvider"]) + start_time = time.time() + onnx_output = session.run(None, {"input": input_tensor.numpy()}) + inference_time = time.time() - start_time + + # Verify accuracy + np.testing.assert_allclose(pytorch_output, onnx_output[0], rtol=1e-3, atol=1e-5) + print("ONNX Runtime verification passed") + return inference_time + + +# Define a function to benchmark OpenVINO inference and verify accuracy +def benchmark_openvino_inference_and_verify(model_path, input_tensor, pytorch_output): + ie = Core() + model_onnx = ie.read_model(model=model_path) + compiled_model = ie.compile_model(model=model_onnx, device_name="CPU") + output_layer = compiled_model.output(0) + + start_time = time.time() + openvino_output = compiled_model([input_tensor.numpy()])[output_layer] + inference_time = time.time() - start_time + + # Verify accuracy + np.testing.assert_allclose(pytorch_output, openvino_output, rtol=1e-3, atol=1e-5) + print("OpenVINO Runtime verification passed") + return inference_time + + +torch.onnx.export(model, + torch.randn(1, 3, 1024, 1024), + 'db_resnet50.onnx', + export_params=True, opset_version=11, do_constant_folding=True, - input_names = ["input"], - output_names = ["output"], - dynamic_axes = {"input":{0:"batch_size", 2:"x_axis", 3:"y_axis"}, - "output":{0:"batch_size", 2:"x_axis", 3:"y_axis"}}) - -import onnxruntime + input_names=["input"], + output_names=["output"], + dynamic_axes={"input": {0: "batch_size", 2: "height", 3: "width"}, + "output": {0: "batch_size", 2: "height", 3: "width"}}) -ort_session = onnxruntime.InferenceSession('db_resnet50_rotation.onnx', providers = ["CPUExecutionProvider"]) +# Example input tensor +input_tensor_1024 = torch.randn(1, 3, 1024, 1024) +input_tensor_1536 = torch.randn(1, 3, 1536, 1536) -ort_inputs = {"input":input.numpy()} +# Perform PyTorch inference and capture output start = time.time() -ort_outs = ort_session.run(None, ort_inputs) -print("onnx time", time.time() - start) -print(np.testing.assert_allclose(pred.detach().cpu().numpy(), ort_outs[0], rtol=1e-3, atol=1e-5)) +pytorch_output_1024 = pytorch_inference(model, input_tensor_1024) +print(f"PyTorch Inference Time (1024x1024): {time.time() - start}") -ort_inputs = {"input":input2.numpy()} start = time.time() -ort_outs = ort_session.run(None, ort_inputs) -print("onnx time", time.time() - start) -start = time.time() -pred = model.det_predictor.model(input2) -print("pytorch time", time.time() - start) -print(np.testing.assert_allclose(pred.detach().cpu().numpy(), ort_outs[0], rtol=1e-3, atol=1e-5)) +pytorch_output_1536 = pytorch_inference(model, input_tensor_1536) +print(f"PyTorch Inference Time (1536x1536): {time.time() - start}") -from openvino.runtime import Core +# Benchmark and verify ONNX Runtime +time_onnx_1024 = benchmark_onnx_inference_and_verify('db_resnet50.onnx', input_tensor_1024, + pytorch_output_1024) +print(f"ONNX Runtime Inference Time (1024x1024): {time_onnx_1024}") -ie = Core() -model_onnx = ie.read_model(model='db_resnet50_rotation.onnx') -start = time.time() -compiled_model_onnx = ie.compile_model(model=model_onnx, device_name="CPU") +time_onnx_1536 = benchmark_onnx_inference_and_verify('db_resnet50.onnx', input_tensor_1536, + pytorch_output_1536) +print(f"ONNX Runtime Inference Time (1536x1536): {time_onnx_1536}") -output_layer_onnx = compiled_model_onnx.output(0) -print("model compilation time", time.time() - start) -start = time.time() -# Run inference on the input image. -print(input2.numpy().shape, input2.dtype) -res_onnx = compiled_model_onnx([input2.numpy()])[output_layer_onnx] -print(res_onnx.shape) -print("openvino time", time.time() - start) -print(np.testing.assert_allclose(pred.detach().cpu().numpy(), res_onnx, rtol=1e-3, atol=1e-5)) +# Benchmark and verify OpenVINO +time_openvino_1024 = benchmark_openvino_inference_and_verify('db_resnet50.onnx', input_tensor_1024, + pytorch_output_1024) +print(f"OpenVINO Inference Time (1024x1024): {time_openvino_1024}") + +time_openvino_1536 = benchmark_openvino_inference_and_verify('db_resnet50.onnx', input_tensor_1536, + pytorch_output_1536) +print(f"OpenVINO Inference Time (1536x1536): {time_openvino_1536}")