feat: add onnx benchmark

Author: naidenovaleksei

config/benchmark_onnx.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+
+---
+defaults:
+  - export: null
+  - _self_
+
+data_dir: ${oc.env:HOME}/data/Niro122-HQ/2023-05-25--09-34-14
+start_frame: 2910
+num_episodes: 50
+frame_step: 10
+output: /tmp/rmind.csv
+warmup: 1
+image_size: null
+
+# override per run, e.g.:
+#   onnx=/path/to/model.onnx wandb_model=yaak/rmind/model-XXXXXXXX:vN
+#   export=yaak/control_transformer/finetuned  # apples-to-apples vs the ONNX export
+onnx: null
+wandb_model: null

justfile

@@ -87,6 +87,13 @@ test *ARGS: generate-config
 update-snapshots:
     uv run python -m tests.scripts.update_snapshots
 
+benchmark-onnx *ARGS: generate-config
+    LD_LIBRARY_PATH="$(find /nix/store -maxdepth 1 -name '*gcc-15*-lib' -type d -print -quit 2>/dev/null)/lib:${LD_LIBRARY_PATH:-}" \
+    uv run --group benchmark rmind-benchmark-onnx \
+        --config-path {{ justfile_directory() }}/config \
+        --config-name benchmark_onnx.yaml \
+        {{ ARGS }}
+
 export-onnx *ARGS: generate-config
     uv run \
         --extra export \

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
 repo = "https://github.com/yaak-ai/rmind"
 
 [project.scripts]
+rmind-benchmark-onnx = "rmind.scripts.benchmark_onnx:main"
 rmind-check-git = "rmind.scripts.check_git:main"
 rmind-export-onnx = "rmind.scripts.export_onnx:main"
 rmind-predict = "rmind.scripts.predict:main"
@@ -76,6 +77,13 @@ dev = [
   "pudb>=2025.1.5",
   "wat-inspector",
 ]
+benchmark = [
+  "onnxruntime>=1.23.1",
+  "opencv-python-headless",
+  "pyproj>=3.6.0",
+  "tabulate>=0.9.0",
+  "wandb>=0.24.2",
+]
 check = [
   "ruff>=0.15.18",
   "ty>=0.0.52",
@@ -104,13 +112,18 @@ DEP002 = [
 ]
 DEP003 = ["rmind"]
 DEP004 = [
+  "cv2",
   "onnxruntime",
+  "pyproj",
+  "tabulate",
   "tensorrt",
 ]
 
 [tool.deptry.package_module_name_map]
 hydra-core = "hydra"
+opencv-python-headless = "cv2"
 tensordict-nightly = "tensordict"
+tensorrt-cu13 = "tensorrt"
 pytorch-lightning = [
   "pytorch_lightning",
   "lightning_fabric",