From 7eca2c133f1ddce9f17f8b0ced63d42499b7804f Mon Sep 17 00:00:00 2001 From: Artur Ayrapetyan Date: Sun, 24 May 2026 12:12:25 +0200 Subject: [PATCH] exported to onnx the pytorch model, at inference load onnx from hf, added benchmark wrt pytorch model, confirmed onnx gains 2.5 speed --- .gitignore | 1 + model/detection.py | 2 +- model/ssd.py | 31 ++++++++---- readme_parallel.md | 76 ++++++++++++++++++++++++++++ requirements/requirements-api.txt | 3 ++ scripts/benchmark_onnx_vs_pytorch.py | 71 ++++++++++++++++++++++++++ scripts/export_onnx.py | 69 +++++++++++++++++++++++++ scripts/test_onnx_image.py | 61 ++++++++++++++++++++++ serving/inference.py | 65 ++++++++++++++++++++++-- serving/server.py | 21 ++++---- 10 files changed, 374 insertions(+), 26 deletions(-) create mode 100644 readme_parallel.md create mode 100644 scripts/benchmark_onnx_vs_pytorch.py create mode 100644 scripts/export_onnx.py create mode 100644 scripts/test_onnx_image.py diff --git a/.gitignore b/.gitignore index 90022d0..ecdfd46 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ labels/ JPEGImages/ **/[Ii]mages test.ipynb +*.onnx # C extensions $.pth *.so diff --git a/model/detection.py b/model/detection.py index 35b2556..786b3bd 100644 --- a/model/detection.py +++ b/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): diff --git a/model/ssd.py b/model/ssd.py index d420b16..cdf17a6 100644 --- a/model/ssd.py +++ b/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) diff --git a/readme_parallel.md b/readme_parallel.md new file mode 100644 index 0000000..7a6f3cf --- /dev/null +++ b/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. + diff --git a/requirements/requirements-api.txt b/requirements/requirements-api.txt index bf9a13b..c81cb38 100644 --- a/requirements/requirements-api.txt +++ b/requirements/requirements-api.txt @@ -11,6 +11,9 @@ opencv-python-headless # Modèle & poids huggingface_hub pyyaml +onnx +onnxscript +onnxruntime # API fastapi diff --git a/scripts/benchmark_onnx_vs_pytorch.py b/scripts/benchmark_onnx_vs_pytorch.py new file mode 100644 index 0000000..6541e82 --- /dev/null +++ b/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() diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py new file mode 100644 index 0000000..5ce49bf --- /dev/null +++ b/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")) diff --git a/scripts/test_onnx_image.py b/scripts/test_onnx_image.py new file mode 100644 index 0000000..99e4b00 --- /dev/null +++ b/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() diff --git a/serving/inference.py b/serving/inference.py index f896133..0a6cd6b 100644 --- a/serving/inference.py +++ b/serving/inference.py @@ -1,5 +1,6 @@ -# inference.py — Import-safe model loading and per-frame detection for the server. -# Separates model setup from CLI/argparse concerns in SsdFastCrowdVision.py. +# defining and loading models from hugging face : using pytorch and onnx (u will have load_ssd_model using pytorch model ) +# qnd load_ssd_model_onnx loading of onnx model , benchmarking (time) is done in scripts\benchmark_onnx_vs_pytorch.py + import torch import numpy as np @@ -14,7 +15,7 @@ import time import logging -logger=logging.getLogger(__name__) +logger = logging.getLogger(__name__) # project root = folder containing this file project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -71,6 +72,64 @@ def load_ssd_model(device): return model, config, transform +def load_ssd_model_onnx(): + """Download ONNX model from HuggingFace and load it for inference. + + Returns (session, config_dict, transform). + NMS is already inside the ONNX graph (NonMaxSuppression op). + """ + import onnxruntime as ort + + onnx_path = hf_hub_download( + repo_id="aayrapet/SsdFastCrowdVision", + filename="SSD_FastCrowdVision_v1.onnx", + ) + logger.info("ONNX model loaded from %s", onnx_path) + + providers = ( + ["CUDAExecutionProvider", "CPUExecutionProvider"] + if torch.cuda.is_available() + else ["CPUExecutionProvider"] + ) + session = ort.InferenceSession(onnx_path, providers=providers) + + wider_yaml = os.path.join(project_root, "datasets", "WiderPeople", "widerpeople.yaml") + with open(wider_yaml) as f: + config = yaml.safe_load(f)["names"] + + transform = v2.Compose(test_val_transform) + + return session, config, transform + + +def detect_frame_onnx(session, pil_image, transform, score_thr=0.25): + """Run the ONNX SSD on a single PIL image. + + NMS runs inside the ONNX graph. Returns np.ndarray of shape (N, 6) — + columns: [x1, y1, x2, y2, score, class] in pixel coordinates. + """ + t0 = time.perf_counter() + W, H = pil_image.size + + x = transform(pil_image).unsqueeze(0).numpy() + topk = session.run(["detections"], {"image": x})[0] + topk = topk.squeeze(0) + + mask = topk[:, 4] > score_thr + if not mask.any(): + return np.empty((0, 6)) + + topk = topk[mask] + + topk[:, [0, 2]] *= W + topk[:, [1, 3]] *= H + + elapsed = time.perf_counter() - t0 + logger.info("[onnx-inference] %.1f ms", elapsed * 1000) + + return topk + + def detect_frame(model, pil_image, transform, device, score_thr=0.25): """Run SSD on a single PIL image. diff --git a/serving/server.py b/serving/server.py index 57179a1..8287559 100644 --- a/serving/server.py +++ b/serving/server.py @@ -12,25 +12,23 @@ import cv2 import numpy as np -import torch from PIL import Image from fastapi import FastAPI, WebSocket, UploadFile, File from fastapi.staticfiles import StaticFiles from norfair import Tracker, Detection -from serving.inference import load_ssd_model, detect_frame +from serving.inference import load_ssd_model_onnx, detect_frame_onnx logger = logging.getLogger(__name__) app = FastAPI() @app.get("/health") def health(): """Endpoint utilisé par le HEALTHCHECK Docker et les sondes Kubernetes.""" - return {"status": "ok", "model_loaded": model is not None} + return {"status": "ok", "model_loaded": session is not None} # --- globals: loaded once at startup, shared across all requests --- -model = None +session = None config = None transform = None -device = None # maps session_id → path of the uploaded temp video file video_sessions: dict[str, str] = {} @@ -40,12 +38,11 @@ def health(): @app.on_event("startup") def startup(): - """Called once when uvicorn starts. Loads the SSD model into memory.""" - global model, config, transform, device - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model, config, transform = load_ssd_model(device) - - logger.info(f"Server ready — model loaded on {device}") + """Called once when uvicorn starts. Loads the ONNX SSD model into memory.""" + global session, config, transform + session, config, transform = load_ssd_model_onnx() + + logger.info("Server ready — ONNX model loaded") @@ -134,7 +131,7 @@ async def detect_ws(websocket: WebSocket): pil_image = Image.fromarray(frame_rgb) # run SSD detection on this frame - detections_np = detect_frame(model, pil_image, transform, device, score_thr) + detections_np = detect_frame_onnx(session, pil_image, transform, score_thr) # convert SSD detections to norfair Detection objects norfair_dets = []