diff --git a/docs/deepstream_validation.md b/docs/deepstream_validation.md index f795b62b..16dd6161 100644 --- a/docs/deepstream_validation.md +++ b/docs/deepstream_validation.md @@ -69,7 +69,7 @@ The models are **validated and ready for performance benchmarking**. - **ONNX predictions:** `/workspace/assets/predictions/onnx/` - **Configs used:** `configs/resnet18.txt`, `configs/mobilenet.txt` - **Engines:** `/workspace/models/deepstream/` -- **Script used:** `src/validate_deepstream_vs_onnx.py` +- **Script used:** `src/model_conversion/validate_trt_vs_onnx.py` - **Execution environments:** - `run_dev.sh` (x86_64) → ONNX inference - `run_dev_jetson.sh` (ARM Jetson) → DeepStream inference diff --git a/docs/export_to_onnx.md b/docs/export_to_onnx.md index 5532de8f..1c1911a4 100644 --- a/docs/export_to_onnx.md +++ b/docs/export_to_onnx.md @@ -8,10 +8,11 @@ All conversions should be executed **inside the dev container** to ensure depend ### Commands Export ResNet18: ```bash -python scripts/export_to_onnx.py resnet18 +python src/model_conversion/pytorch_to_onnx.py resnet18 + ``` Export MobileNetV2: ```bash -python scripts/export_to_onnx.py mobilenet +python src/model_conversion/pytorch_to_onnx.py mobilenet ``` \ No newline at end of file diff --git a/src/deepstream_pipeline.py b/src/deepstream_pipeline.py index b96450a8..7e162a2f 100644 --- a/src/deepstream_pipeline.py +++ b/src/deepstream_pipeline.py @@ -1,12 +1,16 @@ - -import gi - -gi.require_version("Gst", "1.0") import os +import sys from datetime import datetime +import gi from gi.repository import GLib, Gst -from engine_helper import ensure_engine_exists + +# Add workspace root to sys.path +sys.path.append("/workspace") + +from src.model_conversion.onnx_to_trt import build_engine_if_missing + +gi.require_version("Gst", "1.0") rtsp_port = os.environ.get("RTSP_PORT", "8554") @@ -133,7 +137,7 @@ def on_message(bus, msg): def main(): - ensure_engine_exists(CONFIG_FILE) + build_engine_if_missing(CONFIG_FILE) # OS-level setup Gst.init(None) os.makedirs(OUTPUT_DIR, exist_ok=True) diff --git a/src/engine_helper.py b/src/model_conversion/onnx_to_trt.py similarity index 97% rename from src/engine_helper.py rename to src/model_conversion/onnx_to_trt.py index 3f210b1c..1cd870eb 100644 --- a/src/engine_helper.py +++ b/src/model_conversion/onnx_to_trt.py @@ -10,7 +10,7 @@ def file_sha256(path: str) -> str: return hashlib.sha256(f.read()).hexdigest() -def ensure_engine_exists(config_path: str) -> str: +def build_engine_if_missing(config_path: str) -> str: """ Ensure that a TensorRT engine file exists based on a DeepStream config file. If not found, create it using trtexec. diff --git a/src/model_conversion/pytorch_to_onnx.py b/src/model_conversion/pytorch_to_onnx.py new file mode 100644 index 00000000..8696fdd2 --- /dev/null +++ b/src/model_conversion/pytorch_to_onnx.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +import argparse +from pathlib import Path +import subprocess + +import torch + +from src.inference import MobileNetInference, ResNetInference +from src.path_utils import ensure_clean_directory + +NUM_CLASSES = 83 +INPUT_SIZE = (256, 256) +DEVICE = "cpu" +MODELS_DIR_PATH = Path("models") + + +def get_model(model_name: str): + if model_name == "resnet18": + return ResNetInference( + device=DEVICE, + weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt", + num_classes=NUM_CLASSES, + ) + elif model_name == "mobilenet": + return MobileNetInference( + device=DEVICE, + weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt", + num_classes=NUM_CLASSES, + ) + else: + raise ValueError(f"Unsupported model: {model_name}") + + +def main(model_name: str): + pytorch_model = get_model(model_name) + + # Detect parameter dtype (fp16/fp32) and match input accordingly + param_dtype = next( + (p.dtype for p in pytorch_model.model.parameters() if p.is_floating_point()), + torch.float32, + ) + + # Fixed input size: 256x256, 3 channels, batch=1 + dummy = torch.zeros(1, 3, INPUT_SIZE[0], INPUT_SIZE[1], dtype=param_dtype) + + # ONNX path/name + onnx_dir = MODELS_DIR_PATH / "onnx" + ensure_clean_directory(onnx_dir) + onnx_path = onnx_dir / f"{model_name}.onnx" + dynamic_axes = {"input": {0: "batch"}, "output": {0: "batch"}} + + with torch.inference_mode(): + torch.onnx.export( + pytorch_model.model, + dummy, + onnx_path.as_posix(), + input_names=["input"], + output_names=["output"], + dynamic_axes=dynamic_axes, + opset_version=17, + do_constant_folding=True, + ) + + subprocess.run(["python3", "src/model_conversion/validate_onnx_vs_pytorch.py", model_name], check=True) + + print(f"Exported: {onnx_path}") + + +if __name__ == "__main__": + p = argparse.ArgumentParser("Export PyTorch model to ONNX (fixed 256x256 input)") + p.add_argument( + "model_name", + type=str, + choices=["resnet18", "mobilenet"], + help="Name of the model to export (must exist in models directory)", + ) + args = p.parse_args() + + main(args.model_name) diff --git a/src/model_conversion/validate_onnx_vs_pytorch.py b/src/model_conversion/validate_onnx_vs_pytorch.py new file mode 100644 index 00000000..c0adf1f9 --- /dev/null +++ b/src/model_conversion/validate_onnx_vs_pytorch.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +import argparse +import os +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from src.inference import ID_TO_NAME, ClassifierInferenceBase, MobileNetInference, ResNetInference +from src.onnx_model import OnnxClassifierInferenceBase + +NUM_CLASSES = 83 +INPUT_SIZE = (256, 256) + +MODELS_DIR_PATH = Path("models") +IMAGES_DIR_PATH = Path("assets/test_images") +DEVICE = "cpu" +INPUT_SHAPE = (1, 3, *INPUT_SIZE) + + +def assert_the_same_shapes( + model_name: str, + onnx_clf: OnnxClassifierInferenceBase, +): + onnx_inputs = [input.shape for input in onnx_clf.session.get_inputs()] + onnx_outputs = [output.shape for output in onnx_clf.session.get_outputs()] + + assert onnx_inputs == [["batch", 3, *INPUT_SIZE]], "ONNX input shape mismatch" + assert onnx_outputs == [ + ["batch", NUM_CLASSES] + ], f"Number of outputs mismatch: {onnx_outputs} is not [['batch', {NUM_CLASSES}]]" + + +def assert_numerical_accuracy( + model_name: str, + torch_clf: ClassifierInferenceBase, + onnx_clf: OnnxClassifierInferenceBase, + rtol: float = 1e-3, + atol: float = 1e-5, +): + with torch.inference_mode(): + param_dtype = next((p.dtype for p in torch_clf.model.parameters() if p.is_floating_point()), torch.float32) + torch_input = torch.randn(*INPUT_SHAPE, dtype=param_dtype) + torch_logits = torch_clf.model(torch_input).to(torch.float32) + torch_np = torch_logits.detach().cpu().numpy() + + onnx_logits = onnx_clf._forward(torch_input.to(torch.float32)) + onnx_np = onnx_logits.detach().cpu().numpy() + + assert np.allclose(torch_np, onnx_np, rtol=rtol, atol=atol) + + +def assert_same_predictions( + model_name: str, + torch_clf: ClassifierInferenceBase, + onnx_clf: OnnxClassifierInferenceBase, +): + torch_results = {} + onnx_results = {} + img_extensions = {".jpg", ".jpeg", ".png", ".bmp"} + + for filename in os.listdir(IMAGES_DIR_PATH): + file_ext = os.path.splitext(filename)[1].lower() + if file_ext in img_extensions: + image_path = os.path.join(IMAGES_DIR_PATH, filename) + image = Image.open(image_path) + + torch_pred = torch_clf.infer(image) + onnx_pred = onnx_clf.infer(image) + + torch_results[filename] = torch_pred[0]["class_id"] + onnx_results[filename] = onnx_pred[0]["class_id"] + + assert torch_results == onnx_results + + +def main(model_name: str): + if model_name == "resnet18": + torch_model = ResNetInference( + device=DEVICE, + weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt", + num_classes=NUM_CLASSES, + strict=True, + ) + elif model_name == "mobilenet": + torch_model = MobileNetInference( + device=DEVICE, + weights_path=MODELS_DIR_PATH / "pytorch" / f"{model_name}.pt", + num_classes=NUM_CLASSES, + strict=True, + ) + else: + raise ValueError(f"Unsupported model: {model_name}") + onnx_model = OnnxClassifierInferenceBase( + device=DEVICE, + weights_path=MODELS_DIR_PATH / "onnx" / f"{model_name}.onnx", + class_mapping=ID_TO_NAME, + topk=1, + ) + + assert_the_same_shapes(model_name, onnx_model) + assert_numerical_accuracy(model_name, torch_model, onnx_model) + assert_same_predictions(model_name, torch_model, onnx_model) + + +if __name__ == "__main__": + p = argparse.ArgumentParser("Validate ONNX model against PyTorch") + p.add_argument( + "model_name", + type=str, + choices=["resnet18", "mobilenet"], + help="Name of the model to validate", + ) + args = p.parse_args() + main(args.model_name) diff --git a/src/validate_deepstream_vs_onnx.py b/src/model_conversion/validate_trt_vs_onnx.py similarity index 100% rename from src/validate_deepstream_vs_onnx.py rename to src/model_conversion/validate_trt_vs_onnx.py diff --git a/src/test_onnx_pipeline.py b/src/test_onnx_pipeline.py index ee9df4fc..7b00fed4 100755 --- a/src/test_onnx_pipeline.py +++ b/src/test_onnx_pipeline.py @@ -8,11 +8,15 @@ gi.require_version("Gst", "1.0") from gi.repository import Gst # noqa: E402 -from engine_helper import ensure_engine_exists + +# Add workspace root to sys.path +sys.path.append("/workspace") + +from src.model_conversion.onnx_to_trt import build_engine_if_missing def main(config_file): # Create deepstream_engines directory with full permissions - ensure_engine_exists(config_file) + build_engine_if_missing(config_file) Gst.init(None) diff --git a/src/test_onnx_pipeline_real.py b/src/test_onnx_pipeline_real.py index 2a385336..b98fd71c 100644 --- a/src/test_onnx_pipeline_real.py +++ b/src/test_onnx_pipeline_real.py @@ -7,10 +7,14 @@ gi.require_version("Gst", "1.0") from gi.repository import GLib, Gst # noqa: E402 -from engine_helper import ensure_engine_exists + +# Add workspace root to sys.path +sys.path.append("/workspace") + +from src.model_conversion.onnx_to_trt import build_engine_if_missing def main(config_file, input_path, is_video=True): - ensure_engine_exists(config_file) + build_engine_if_missing(config_file) Gst.init(None) pipeline = Gst.Pipeline()