Rust tensorrt warmup script (HULKs#2312)

* Add rust script to precompile onnx models

* Run pepsi format

* Fix image width

Author: oleflb

Cargo.lock

@@ -57,6 +57,7 @@ members = [
   "tools/mujoco-simulator/mujoco-rust-server",
   "tools/parameter_tester",
   "tools/pepsi",
+  "tools/tensorrt-compile",
   "tools/twix",
   "tools/vista",
   "tools/widget_gallery",

Cargo.toml

@@ -5,10 +5,10 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
-    "click>=8.3.1",
-    "onnxruntime>=1.24.3",
-    "onnxscript>=0.6.2",
-    "onnxslim>=0.1.87",
-    "torch>=2.10.0",
-    "ultralytics>=8.4.21",
+  "click>=8.3.1",
+  "onnxruntime>=1.24.3",
+  "onnxscript>=0.6.2",
+  "onnxslim>=0.1.87",
+  "torch>=2.10.0",
+  "ultralytics>=8.4.21",
 ]

tools/machine-learning/multi-task-yolo/pyproject.toml

@@ -0,0 +1,12 @@
+[package]
+name = "tensorrt-compile"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+homepage.workspace = true
+
+[dependencies]
+clap = { workspace = true }
+color-eyre = { workspace = true }
+ndarray = { workspace = true }
+ort = { workspace = true, features = ["cuda", "tensorrt"] }

tools/tensorrt-compile/Cargo.toml

@@ -0,0 +1,56 @@
+use std::path::PathBuf;
+
+use clap::Parser;
+use color_eyre::{Result, eyre::Context};
+use ndarray::Array3;
+use ort::{
+    execution_providers::{CUDAExecutionProvider, TensorRTExecutionProvider},
+    inputs,
+    session::{Session, SessionOutputs, builder::GraphOptimizationLevel},
+    value::TensorRef,
+};
+
+#[derive(Debug, Parser)]
+struct CliArguments {
+    /// Path to onnx model
+    onnx_path: PathBuf,
+
+    /// Path to cache folder
+    #[arg(long, default_value = "/home/booster/.cache/hulk/tensor-rt")]
+    cache_path: PathBuf,
+}
+
+fn main() -> Result<()> {
+    const IMAGE_WIDTH: usize = 544;
+    const IMAGE_HEIGHT: usize = 448;
+
+    let args = CliArguments::parse();
+    color_eyre::install()?;
+    std::fs::create_dir_all(&args.cache_path).wrap_err("failed to create cache path")?;
+
+    let tensor_rt = TensorRTExecutionProvider::default()
+        .with_device_id(0)
+        .with_fp16(true)
+        .with_engine_cache(true)
+        .with_engine_cache_path(args.cache_path.display())
+        .build()
+        .error_on_failure();
+    let cuda = CUDAExecutionProvider::default().build();
+
+    let mut session = Session::builder()?
+        .with_execution_providers([tensor_rt, cuda])?
+        .with_optimization_level(GraphOptimizationLevel::Level3)?
+        .with_intra_threads(2)?
+        .commit_from_file(args.onnx_path)?;
+
+    let sample_image = Array3::<u8>::default([IMAGE_HEIGHT / 2, IMAGE_WIDTH / 2, 6]);
+    let outputs: SessionOutputs = session
+        .run(inputs!["raw_bytes_input" => TensorRef::from_array_view(sample_image.view())?])?;
+    let _ = outputs["network_detections"]
+        .try_extract_array::<f32>()?
+        .t()
+        .into_owned();
+    eprintln!("object detection setup complete");
+
+    Ok(())
+}