116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
import os
|
|
import time
|
|
import argparse
|
|
import torch
|
|
import torchvision
|
|
import onnx
|
|
import onnxruntime as ort
|
|
import numpy as np
|
|
|
|
def export_to_onnx(model_path, onnx_path):
|
|
print(f"Loading PyTorch checkpoint from {model_path}...")
|
|
from data_generator import CLASSES
|
|
num_classes = len(CLASSES) + 1
|
|
|
|
# Recreate the model structure
|
|
model = torchvision.models.detection.ssdlite320_mobilenet_v3_large(
|
|
num_classes=num_classes,
|
|
weights_backbone=torchvision.models.MobileNet_V3_Large_Weights.DEFAULT
|
|
)
|
|
|
|
# Load state dict
|
|
state_dict = torch.load(model_path, map_location='cpu')
|
|
model.load_state_dict(state_dict)
|
|
model.eval()
|
|
|
|
# Dummy input: list of 1 tensor of size [3, 320, 320]
|
|
# SSDLite in torchvision expects a list of tensors during export
|
|
dummy_input = [torch.randn(3, 320, 320)]
|
|
|
|
print("Exporting model to ONNX...")
|
|
os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
|
|
|
|
torch.onnx.export(
|
|
model,
|
|
(dummy_input,),
|
|
onnx_path,
|
|
opset_version=12,
|
|
input_names=["images"],
|
|
output_names=["boxes", "scores", "labels"],
|
|
dynamic_axes={
|
|
"images": {0: "batch_size"}
|
|
},
|
|
dynamo=False
|
|
)
|
|
|
|
print(f"Successfully saved ONNX model to {onnx_path}")
|
|
|
|
# Verify ONNX model
|
|
print("Verifying ONNX model graph structure...")
|
|
onnx_model = onnx.load(onnx_path)
|
|
onnx.checker.check_model(onnx_model)
|
|
|
|
print("Inputs:")
|
|
for input_node in onnx_model.graph.input:
|
|
print(f" - Name: {input_node.name}")
|
|
print("Outputs:")
|
|
for output_node in onnx_model.graph.output:
|
|
print(f" - Name: {output_node.name}")
|
|
|
|
def benchmark_onnx(onnx_path):
|
|
print(f"Benchmarking ONNX inference speed using ONNX Runtime...")
|
|
|
|
# Create inference session
|
|
providers = ['CPUExecutionProvider']
|
|
if 'DirectMLExecutionProvider' in ort.get_available_providers():
|
|
providers = ['DirectMLExecutionProvider', 'CPUExecutionProvider']
|
|
print("Using DirectML hardware acceleration!")
|
|
|
|
session = ort.InferenceSession(onnx_path, providers=providers)
|
|
|
|
# Prepare dummy input matching the shape of the exported model input
|
|
# Note: torchvision SSD ONNX export expects list of images.
|
|
# The input node shape in ONNX has dynamic axes. Let's create a float32 array
|
|
# corresponding to the image. SSDLite internally resizes and batches.
|
|
# Let's see the input signature. The input 'images' expects a 3D or 4D float tensor.
|
|
# In our export script, the dummy input was [torch.randn(3, 320, 320)]
|
|
# In ONNX, a list of Tensors gets mapped to a single concatenated batch tensor or similar.
|
|
# Let's generate a random array of shape (1, 3, 320, 320)
|
|
input_name = session.get_inputs()[0].name
|
|
dummy_in = np.random.randn(3, 320, 320).astype(np.float32)
|
|
|
|
# Warmup
|
|
for _ in range(5):
|
|
_ = session.run(None, {input_name: dummy_in})
|
|
|
|
# Measure latency
|
|
num_runs = 50
|
|
start_time = time.time()
|
|
for _ in range(num_runs):
|
|
_ = session.run(None, {input_name: dummy_in})
|
|
elapsed = time.time() - start_time
|
|
avg_latency_ms = (elapsed / num_runs) * 1000
|
|
|
|
print(f"Average ONNX Runtime inference latency: {avg_latency_ms:.2f} ms")
|
|
if avg_latency_ms < 50.0:
|
|
print("Success: Inference latency is under the 50 ms limit!")
|
|
else:
|
|
print("Warning: Inference latency is above the 50 ms limit. Optimize hardware providers.")
|
|
|
|
return avg_latency_ms
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Export PyTorch SSDLite model to ONNX")
|
|
parser.add_argument("--pytorch-model", type=str, default=os.path.join("models", "model.pt"), help="Path to PyTorch model weights")
|
|
parser.add_argument("--onnx-model", type=str, default=os.path.join("models", "model.onnx"), help="Path to save ONNX model")
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
export_to_onnx(args.pytorch_model, args.onnx_model)
|
|
benchmark_onnx(args.onnx_model)
|
|
except Exception as e:
|
|
print(f"Error during export: {e}")
|
|
import sys
|
|
sys.exit(1)
|