onnx export

.gitignore

@@ -13,6 +13,7 @@ labels/
 JPEGImages/
 **/[Ii]mages
 test.ipynb
+*.onnx
 # C extensions
 $.pth
 *.so

model/detection.py

@@ -6,7 +6,7 @@
 
 class Detection(nn.Module):
     """
-    Performing Non-max Supression on model outputs 
+    Performing Non-max Supression on model outputs thanks to https://github.com/amdegroot/ssd.pytorch
 
     """
     def __init__(self, nb_classes, prob_thr, nms_thr, top_k, variances: list, anchors):

model/ssd.py

@@ -148,11 +148,11 @@ def __init__(
     def _hook_fn(self, module, input, output):
         self._hooked_features.append(output)
 
-    def forward(self, X):
-        self._hooked_features = []  
+    def _compute_locs_confs(self, X):
+        """Shared feature extraction used by forward() and ONNX export."""
+        self._hooked_features = []
         layers_for_prediction = []
 
-        # base model
         c5 = self.backbone(X)
         c4 = self._hooked_features[0]
         if self.c4_norm is not None:
@@ -162,14 +162,12 @@ def forward(self, X):
 
         for idx in range(len(self.extras)):
             X = self.extras[idx](X)
-
             layers_for_prediction.append(X)
 
         classifications = []
         for layer_for_predictions, classification_convolution in zip(
             layers_for_prediction, self.classification_convolutions
         ):
-
             x = classification_convolution(layer_for_predictions)
             # then we want to get for all i,j in H*H and all k in 1....K -> p1.....pC probabilities of C classes
             """
@@ -199,26 +197,39 @@ def forward(self, X):
             regressions.append(x.permute(0, 2, 3, 1).contiguous())
 
         # this efficient code was taken from degroot/ssd.pytorch github and is equivalent to my code in comment
-
         loc = torch.cat([o.view(o.size(0), -1) for o in regressions], 1)
         conf = torch.cat([o.view(o.size(0), -1) for o in classifications], 1)
 
-        locs = loc.view(
-            loc.size(0), -1, 4
-        )  # so we get for every image 2D matrix for classification and regression : 8732 anchor boxes and nb coords/classes
+        locs = loc.view(loc.size(0), -1, 4)
         # 8732 anchor boxes are sum of all anchor boxes across all ft map k =  of sum over k (Hk*Hk*ak)
         # for standard ssd300 it is 38*38*4+19*19*6+100*6+25*6+9*4+4
         confs = conf.view(conf.size(0), -1, self.nb_classes)
+        return locs, confs
+
+    def forward(self, X):
+        locs, confs = self._compute_locs_confs(X)
 
         if self.phase == "train":
             return locs, confs
         elif self.phase == "test":
             output = self.detection(confs, locs)
-            return locs, confs,output
+            return locs, confs, output
         else:
             raise ValueError("Unknown phase. Expected train or test ")
 
 
+class SSDOnnxWrapper(nn.Module):
+    """Full inference graph for ONNX export: SSD features + Detection (NMS)."""
+
+    def __init__(self, ssd: SSD):
+        super().__init__()
+        self.ssd = ssd
+
+    def forward(self, x):
+        locs, confs = self.ssd._compute_locs_confs(x)
+        return self.ssd.detection(confs, locs)
+
+
 def xavier(param):
     init.xavier_uniform_(param)
 

readme_parallel.md

@@ -0,0 +1,76 @@
+# ONNX export 
+
+This project runs object detection algorithm Single Shot Detection using Lightweight Neural Networks such as MobilenetV3 for backbone (2019), the objective is to efficiently track people in dense city areas on videos on CPUs, for this I use detection over frames and simple SORT tracking algorithm using Kalman Filter. For model references, deployment and overall vision, please refer to `README.md`. For now, i use only web for video processing, but in future the objective is to run these algorithms on edge devices, such as mobile phones or cheap Raspberry chips. 
+
+For training, the Wider People Dataset was used for training, next the trained model was uploaded to Hugging Face (https://huggingface.co/aayrapet).
+
+For inference the original project ran  entirely in PyTorch: the server built an SSDLite model, downloaded `SSD_FastCrowdVision_v1.pth` from HuggingFace, and applied post-processing (NMS) through the Python `Detection` module in `model/detection.py`. To speed up I export the model to ONNX. The export script loads the same trained weights from HuggingFace, wraps the network in `SSDOnnxWrapper` (defined in `model/ssd.py`) that is then converted to onnx format. During export, NMS from `model/detection.py` become ONNX's built-in `NonMaxSuppression` operator (https://onnx.ai/onnx/operators/onnx__NonMaxSuppression.html).
+
+
+On the inference side, `serving/inference.py` now exposes two loaders. `load_ssd_model(device)` is the original PyTorch path, which was used before to do forward pass for image. I added `load_ssd_model_onnx()` that downloads `SSD_FastCrowdVision_v1.onnx` from the same HuggingFace repo (`aayrapet/SsdFastCrowdVision`) and creates an ONNX Runtime session; NMS is already inside the graph. `detect_frame_onnx()` mirrors `detect_frame()` so the rest of the pipeline (tracking, drawing) stays the same. The FastAPI server in `serving/server.py` was switched to the ONNX loader and `detect_frame_onnx`, so Docker and local uvicorn both use the HF ONNX model at startup without needing a local `.pth` file.
+
+New scripts under `scripts/` support the workflow: `export_onnx.py` produces the ONNX model, `test_onnx_image.py` runs one image and saves an annotated result, and `benchmark_onnx_vs_pytorch.py` compares latency between the HF PyTorch and HF ONNX models on CPU with a random tensor. 
+
+## Download a test image to the project root
+
+A common WiderPeople test photo is hosted on SSP Cloud MinIO. To copy it into the repository root as `image000047.jpg`, run from the project folder.
+
+On Windows PowerShell:
+
+```powershell
+Invoke-WebRequest -Uri "https://minio.lab.sspcloud.fr/aayrapetyan/FastCrowdVision/datasets/WiderPeople/test/images/image000047.jpg" -OutFile "image000047.jpg"
+```
+
+On Linux or macOS:
+
+```bash
+curl -L -o image000047.jpg "https://minio.lab.sspcloud.fr/aayrapetyan/FastCrowdVision/datasets/WiderPeople/test/images/image000047.jpg"
+```
+
+You can also download the file in a browser and save it manually as `image000047.jpg` in the FastCrowdVision root.
+
+## Export the ONNX model
+
+Export requires  `requirements-api.txt`. Activate your venv.
+
+Then export model to onnx 
+
+```powershell
+python scripts/export_onnx.py
+```
+
+The script downloads `SSD_FastCrowdVision_v1.pth` from HuggingFace, loads weights into SSDLite, wraps the model with `SSDOnnxWrapper`, and writes `SSD_FastCrowdVision_v1.onnx` in the project root. Export uses the legacy ONNX tracer (`dynamo=False`) because the detection post-processing contains Python control flow in detection that the newer exporter does not handle yet. After export, I uploaded `SSD_FastCrowdVision_v1.onnx` to the HuggingFace repo.
+
+## Test ONNX inference on one image
+
+
+```powershell
+python scripts/test_onnx_image.py image000047.jpg
+```
+
+This calls `load_ssd_model_onnx()`, which downloads the ONNX model from HuggingFace, runs inference, and saves `result_onnx.jpg` in the project root with bounding boxes drawn. 
+
+## Benchmark PyTorch vs ONNX
+
+To compare inference speed on CPU with the same random input:
+
+```powershell
+python scripts/benchmark_onnx_vs_pytorch.py
+```
+
+Both backends load from HuggingFace: `.pth` for PyTorch and `.onnx` for ONNX Runtime. The script prints milliseconds per frame and FPS for each, plus a speedup factor relative to PyTorch. Timing excludes the image transform so the comparison focuses on forward pass and NMS.
+
+
+## Docker
+
+same as before (refer more to REAMDE.md for deployment)
+
+```powershell
+docker build -t fastcrowdvision .
+docker run -p 8000:8000 fastcrowdvision
+```
+
+## Conclusion
+
+This exercise allowed me to go further after model development and basic deployment and export the model to ONNX, this is interesting because first ONNX does not require pytorch at all to run in applications and it is faster (confirmed with benchmarks, 2.5 faster). However, for non standard neural networks, as in image detection we use here (when there is data post processing after forward pass as NMS), it is better to use legacy ONNC export.
+

requirements/requirements-api.txt

@@ -11,6 +11,9 @@ opencv-python-headless
 # Modèle & poids
 huggingface_hub
 pyyaml
+onnx
+onnxscript
+onnxruntime
 
 # API
 fastapi

scripts/benchmark_onnx_vs_pytorch.py

@@ -0,0 +1,71 @@
+import os
+import sys
+import time
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from serving.inference import load_ssd_model, load_ssd_model_onnx
+
+N_WARMUP = 10
+N_ITERS = 100
+DEVICE = torch.device("cpu")
+
+
+def _benchmark_pytorch(model, x, n_warmup, n_iters):
+    model.eval()
+    model.phase = "test"
+
+    with torch.no_grad():
+        for _ in range(n_warmup):
+            model(x)
+        t0 = time.perf_counter()
+        for _ in range(n_iters):
+            model(x)
+        elapsed = time.perf_counter() - t0
+
+    return elapsed / n_iters
+
+
+def _benchmark_onnx(session, x_np, n_warmup, n_iters):
+    for _ in range(n_warmup):
+        session.run(["detections"], {"image": x_np})
+
+    t0 = time.perf_counter()
+    for _ in range(n_iters):
+        session.run(["detections"], {"image": x_np})
+    elapsed = time.perf_counter() - t0
+
+    return elapsed / n_iters
+
+
+def _print_result(name, mean_s, baseline_s=None):
+    fps = 1.0 / mean_s
+    ms = mean_s * 1000
+    line = f"{name:8s}  {ms:7.2f} ms/frame  ({fps:6.1f} FPS)"
+    if baseline_s is not None and mean_s > 0:
+        speedup = baseline_s / mean_s
+        line += f"  — {speedup:.2f}x vs PyTorch"
+    print(line)
+
+
+def main():
+    #test on random image 
+    print("Input: random tensor (1, 3, 300, 300)")
+    #load both models from hf, compare .pth model vs onnx model
+    model, _, _ = load_ssd_model(DEVICE)
+    session, _, _ = load_ssd_model_onnx()
+
+    x = torch.randn(1, 3, 300, 300)
+    x_onnx = x.numpy()
+
+    pytorch_ms = _benchmark_pytorch(model, x, N_WARMUP, N_ITERS)
+    onnx_ms = _benchmark_onnx(session, x_onnx, N_WARMUP, N_ITERS)
+
+    _print_result("PyTorch", pytorch_ms)
+    _print_result("ONNX", onnx_ms, baseline_s=pytorch_ms)
+
+
+if __name__ == "__main__":
+    main()

scripts/export_onnx.py

@@ -0,0 +1,69 @@
+import os
+import sys
+
+import torch
+from huggingface_hub import hf_hub_download
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from model.mobilenetv3 import MobileNetV3Large
+from model.ssd import SSDLite, SSDOnnxWrapper
+from training.eval import load_model
+
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def build_loaded_model(device: torch.device):
+    mn = MobileNetV3Large(0.1, 1000, 1280)
+    backbone = mn.features[:-1]
+
+    model = SSDLite(
+        backbone_config_path=os.path.join(project_root, "config", "ssdlite_mobilenetv3large.yaml"),
+        backbone=backbone,
+        c4_name="7.features.1.features.2",
+        nb_classes=4,
+        phase="test",
+        alpha=1.0,
+        prob_thr=0.01,
+        nms_thr=0.45,
+        top_k=200,
+        variances=[0.1, 0.2],
+        N_epochs=50,
+        device=device,
+    ).to(device)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0005)
+
+    weights_path = hf_hub_download(
+        repo_id="aayrapet/SsdFastCrowdVision",
+        filename="SSD_FastCrowdVision_v1.pth",
+    )
+    model, _, _, max_map, _ = load_model(weights_path, device, model, optimizer)
+    model.eval()
+    return model, max_map
+
+
+def export(device: torch.device):
+    output_path = os.path.join(project_root, "SSD_FastCrowdVision_v1.onnx")
+    model, max_map = build_loaded_model(device)
+    wrapper = SSDOnnxWrapper(model).eval()
+
+    print(f"Loaded weights — mAP on WiderPeople: {max_map}")
+
+    # random input ensures NMS path is traced during export
+    dummy = torch.randn(1, 3, 300, 300, device=device)
+
+    torch.onnx.export(
+        wrapper,
+        dummy,
+        output_path,
+        input_names=["image"],
+        output_names=["detections"],
+        opset_version=12,
+        dynamo=False,
+    )
+    print(f"Exported ONNX model (with NMS) to {output_path}")
+
+
+if __name__ == "__main__":
+    export(torch.device("cpu"))

scripts/test_onnx_image.py

@@ -0,0 +1,61 @@
+import argparse
+import os
+import sys
+
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from serving.draw_inference import draw_boxes_with_labels
+from serving.inference import load_ssd_model_onnx
+
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Test ONNX SSD on one image")
+    parser.add_argument("image_path", type=str, help="Path to input image (jpeg/png)")
+    parser.add_argument(
+        "--output",
+        default=os.path.join(project_root, "result_onnx.jpg"),
+        help="Path for annotated output image",
+    )
+    parser.add_argument(
+        "--score_thr",
+        default=0.25,
+        type=float,
+        help="Probability threshold for filtering boxes",
+    )
+    parser.add_argument(
+        "--show_labels",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="Draw class labels and scores on boxes",
+    )
+    args = parser.parse_args()
+
+    session, config, transform = load_ssd_model_onnx()
+
+    img = Image.open(args.image_path).convert("RGB")
+    W, H = img.size
+
+    x = transform(img).unsqueeze(0).numpy()
+    topk = torch.from_numpy(session.run(["detections"], {"image": x})[0])
+    print(f"Detections above threshold: {(topk[0, :, 4] > args.score_thr).sum().item()}")
+
+    out = draw_boxes_with_labels(
+        img,
+        config,
+        topk,
+        H,
+        W,
+        score_thr=args.score_thr,
+        show_scores=args.show_labels,
+    )
+    out.save(args.output)
+    print(f"Saved result to {args.output}")
+
+
+if __name__ == "__main__":
+    main()