initial commit

export.py

@@ -0,0 +1,114 @@
+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"}
+        }
+    )
+    
+    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)