feat(yoloe): support TensorRT/CUDA/CoreML; export ONNX models with dynamic axes (#150)

Author: jamjamjon

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "usls"
 edition = "2021"
-version = "0.1.3"
+version = "0.1.4"
 rust-version = "1.85"
 description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models."
 repository = "https://github.com/jamjamjon/usls"

examples/yoloe-text-prompt/main.rs

@@ -21,14 +21,18 @@ struct Args {
         option,
         default = "vec![
             String::from(\"person\"),
-            String::from(\"bus\"),
             String::from(\"dog\"),
+            String::from(\"bus\"),
             String::from(\"cat\"),
             String::from(\"sign\"),
             String::from(\"tree\"),
         ]"
     )]
     labels: Vec<String>,
+
+    /// batch size
+    #[argh(option, default = "1")]
+    batch_size: usize,
 }
 
 fn main() -> Result<()> {
@@ -39,10 +43,11 @@ fn main() -> Result<()> {
     let args: Args = argh::from_env();
 
     // config
-    let config = Config::yoloe_v8l_seg_tp_80()
+    let config = Config::yoloe_v8m_seg_tp()
+        .with_batch_size_all_min_opt_max(1, args.batch_size, 8)
         .with_model_dtype(args.dtype.as_str().parse()?)
-        .with_model_device(args.device.as_str().parse()?)
-        .with_batch_size_all(1)
+        .with_textual_dtype("fp16".parse()?) // Use FP32 when TensorRT is enabled
+        .with_device_all(args.device.as_str().parse()?)
         .commit()?;
     let mut model = YOLO::new(config)?;
 

src/core/config.rs

@@ -198,6 +198,33 @@ impl Config {
         self
     }
 
+    pub fn with_batch_size_all_min_opt_max(mut self, min: usize, opt: usize, max: usize) -> Self {
+        self.visual = self.visual.with_ixx(0, 0, (min, opt, max).into());
+        self.textual = self.textual.with_ixx(0, 0, (min, opt, max).into());
+        self.model = self.model.with_ixx(0, 0, (min, opt, max).into());
+        self.encoder = self.encoder.with_ixx(0, 0, (min, opt, max).into());
+        self.decoder = self.decoder.with_ixx(0, 0, (min, opt, max).into());
+        self.visual_encoder = self.visual_encoder.with_ixx(0, 0, (min, opt, max).into());
+        self.textual_encoder = self.textual_encoder.with_ixx(0, 0, (min, opt, max).into());
+        self.visual_decoder = self.visual_decoder.with_ixx(0, 0, (min, opt, max).into());
+        self.textual_decoder = self.textual_decoder.with_ixx(0, 0, (min, opt, max).into());
+        self.textual_decoder_merged =
+            self.textual_decoder_merged
+                .with_ixx(0, 0, (min, opt, max).into());
+        self.size_encoder = self.size_encoder.with_ixx(0, 0, (min, opt, max).into());
+        self.size_decoder = self.size_decoder.with_ixx(0, 0, (min, opt, max).into());
+        self.coord_encoder = self.coord_encoder.with_ixx(0, 0, (min, opt, max).into());
+        self.coord_decoder = self.coord_decoder.with_ixx(0, 0, (min, opt, max).into());
+        self.visual_projection = self
+            .visual_projection
+            .with_ixx(0, 0, (min, opt, max).into());
+        self.textual_projection = self
+            .textual_projection
+            .with_ixx(0, 0, (min, opt, max).into());
+
+        self
+    }
+
     pub fn with_batch_size_all(mut self, batch_size: usize) -> Self {
         self.visual = self.visual.with_ixx(0, 0, batch_size.into());
         self.textual = self.textual.with_ixx(0, 0, batch_size.into());

src/core/ort_config.rs

@@ -126,7 +126,6 @@ impl ORTConfig {
         self
     }
 
-    // Hardware configuration methods
     pub fn with_cpu_arena_allocator(mut self, x: bool) -> Self {
         self.hardware.cpu.arena_allocator = x;
         self

src/core/ort_engine.rs

@@ -694,7 +694,7 @@ impl OrtEngine {
                         .with_static_input_shapes(self.hardware.coreml.static_input_shapes)
                         .with_subgraphs(self.hardware.coreml.subgraph_running)
                         .with_compute_units(ort::execution_providers::coreml::CoreMLComputeUnits::All)
-                        .with_model_format(ort::execution_providers::coreml::CoreMLModelFormat::MLProgram)
+                        // .with_model_format(ort::execution_providers::coreml::CoreMLModelFormat::MLProgram)
                         .with_specialization_strategy(
                             ort::execution_providers::coreml::CoreMLSpecializationStrategy::FastPrediction,
                         );

src/models/yolo/config.rs

@@ -180,205 +180,57 @@ impl Config {
 
     fn yoloe_seg_tp() -> Self {
         Self::yolo()
-            .with_num_dry_run_all(0) // TODO: The text encoder must align with the number of classes; running a dry run will cause an error.
             .with_batch_size_all(1)
+            .with_nc(80)
+            .with_model_ixx(1, 1, (1, 80, 300).into()) // max_text_classes
             .with_task(Task::InstanceSegmentation)
             .with_textual_file("mobileclip/blt-textual.onnx")
             .with_model_max_length(77)
+            .with_textual_ixx(0, 1, 77.into())
             .with_tokenizer_file("clip/tokenizer.json")
             .with_tokenizer_config_file("clip/tokenizer_config.json")
             .with_special_tokens_map_file("clip/special_tokens_map.json")
     }
 
-    pub fn yoloe_v8s_seg_tp_80() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::S)
-            .with_nc(80)
-            .with_model_file("yoloe-v8s-seg-tp-80.onnx")
-    }
-
-    pub fn yoloe_v8m_seg_tp_80() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::M)
-            .with_nc(80)
-            .with_model_file("yoloe-v8m-seg-tp-80.onnx")
-    }
-
-    pub fn yoloe_v8l_seg_tp_80() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::L)
-            .with_nc(80)
-            .with_model_file("yoloe-v8l-seg-tp-80.onnx")
-    }
-
-    pub fn yoloe_v8s_seg_tp_300() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::S)
-            .with_nc(300)
-            .with_model_file("yoloe-v8s-seg-tp-300.onnx")
-    }
-
-    pub fn yoloe_v8m_seg_tp_300() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::M)
-            .with_nc(300)
-            .with_model_file("yoloe-v8m-seg-tp-300.onnx")
-    }
-
-    pub fn yoloe_v8l_seg_tp_300() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::L)
-            .with_nc(300)
-            .with_model_file("yoloe-v8l-seg-tp-300.onnx")
-    }
-
-    pub fn yoloe_v8s_seg_tp_500() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::S)
-            .with_nc(500)
-            .with_model_file("yoloe-v8s-seg-tp-500.onnx")
-    }
-
-    pub fn yoloe_v8m_seg_tp_500() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::M)
-            .with_nc(500)
-            .with_model_file("yoloe-v8m-seg-tp-500.onnx")
-    }
-
-    pub fn yoloe_v8l_seg_tp_500() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(8.into())
-            .with_scale(Scale::L)
-            .with_nc(500)
-            .with_model_file("yoloe-v8l-seg-tp-500.onnx")
-    }
-
-    pub fn yoloe_v8s_seg_tp_2000() -> Self {
+    pub fn yoloe_v8s_seg_tp() -> Self {
         Self::yoloe_seg_tp()
             .with_version(8.into())
             .with_scale(Scale::S)
-            .with_nc(2000)
-            .with_model_file("yoloe-v8s-seg-tp-2000.onnx")
+            .with_model_file("yoloe-v8s-seg-tp.onnx")
     }
 
-    pub fn yoloe_v8m_seg_tp_2000() -> Self {
+    pub fn yoloe_v8m_seg_tp() -> Self {
         Self::yoloe_seg_tp()
             .with_version(8.into())
             .with_scale(Scale::M)
-            .with_nc(2000)
-            .with_model_file("yoloe-v8m-seg-tp-2000.onnx")
+            .with_model_file("yoloe-v8m-seg-tp.onnx")
     }
 
-    pub fn yoloe_v8l_seg_tp_2000() -> Self {
+    pub fn yoloe_v8l_seg_tp() -> Self {
         Self::yoloe_seg_tp()
             .with_version(8.into())
             .with_scale(Scale::L)
-            .with_nc(2000)
-            .with_model_file("yoloe-v8l-seg-tp-2000.onnx")
-    }
-
-    pub fn yoloe_11s_seg_tp_80() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::S)
-            .with_nc(80)
-            .with_model_file("yoloe-11s-seg-tp-80.onnx")
-    }
-
-    pub fn yoloe_11m_seg_tp_80() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::M)
-            .with_nc(80)
-            .with_model_file("yoloe-11m-seg-tp-80.onnx")
-    }
-
-    pub fn yoloe_11l_seg_tp_80() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::L)
-            .with_nc(80)
-            .with_model_file("yoloe-11l-seg-tp-80.onnx")
-    }
-
-    pub fn yoloe_11s_seg_tp_300() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::S)
-            .with_nc(300)
-            .with_model_file("yoloe-11s-seg-tp-300.onnx")
-    }
-
-    pub fn yoloe_11m_seg_tp_300() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::M)
-            .with_nc(300)
-            .with_model_file("yoloe-11m-seg-tp-300.onnx")
-    }
-
-    pub fn yoloe_11l_seg_tp_300() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::L)
-            .with_nc(300)
-            .with_model_file("yoloe-11l-seg-tp-300.onnx")
-    }
-
-    pub fn yoloe_11s_seg_tp_500() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::S)
-            .with_nc(500)
-            .with_model_file("yoloe-11s-seg-tp-500.onnx")
-    }
-
-    pub fn yoloe_11m_seg_tp_500() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::M)
-            .with_nc(500)
-            .with_model_file("yoloe-11m-seg-tp-500.onnx")
-    }
-
-    pub fn yoloe_11l_seg_tp_500() -> Self {
-        Self::yoloe_seg_tp()
-            .with_version(11.into())
-            .with_scale(Scale::L)
-            .with_nc(500)
-            .with_model_file("yoloe-11l-seg-tp-500.onnx")
+            .with_model_file("yoloe-v8l-seg-tp.onnx")
     }
 
-    pub fn yoloe_11s_seg_tp_2000() -> Self {
+    pub fn yoloe_11s_seg_tp() -> Self {
         Self::yoloe_seg_tp()
             .with_version(11.into())
             .with_scale(Scale::S)
-            .with_nc(2000)
-            .with_model_file("yoloe-11s-seg-tp-2000.onnx")
+            .with_model_file("yoloe-11s-seg-tp.onnx")
     }
 
-    pub fn yoloe_11m_seg_tp_2000() -> Self {
+    pub fn yoloe_11m_seg_tp() -> Self {
         Self::yoloe_seg_tp()
             .with_version(11.into())
             .with_scale(Scale::M)
-            .with_nc(2000)
-            .with_model_file("yoloe-11m-seg-tp-2000.onnx")
+            .with_model_file("yoloe-11m-seg-tp.onnx")
     }
 
-    pub fn yoloe_11l_seg_tp_2000() -> Self {
+    pub fn yoloe_11l_seg_tp() -> Self {
         Self::yoloe_seg_tp()
             .with_version(11.into())
             .with_scale(Scale::L)
-            .with_nc(2000)
-            .with_model_file("yoloe-11l-seg-tp-2000.onnx")
+            .with_model_file("yoloe-11l-seg-tp.onnx")
     }
 }