Fix export test matrices to exclude nms from Classify models (#18880)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

docs/overrides/javascript/benchmark.js

@@ -116,18 +116,14 @@ function updateChart(initialDatasets = []) {
     EfficientDet: "#000000",
   };
 
-  // Get the selected algorithms from the initialDatasets or all if empty.
+  // Always include all models in the dataset creation
-  const selectedAlgorithms =
+  const datasets = Object.keys(data).map((algorithm, i) => {
-    initialDatasets.length > 0 ? initialDatasets : Object.keys(data);
-
-  // Create the datasets for the selected algorithms.
-  const datasets = selectedAlgorithms.map((algorithm, i) => {
     const baseColor =
       colorMap[algorithm] || `hsl(${Math.random() * 360}, 70%, 50%)`;
     const lineColor =
       Object.keys(data).indexOf(algorithm) === 0
         ? baseColor
-        : lightenHexColor(baseColor, 0.6); // Lighten non-primary lines
+        : lightenHexColor(baseColor, 0.6);
 
     return {
       label: algorithm,
@@ -137,14 +133,15 @@ function updateChart(initialDatasets = []) {
         version: version.toUpperCase(),
       })),
       fill: false,
-      borderColor: lineColor, // Use the lightened color for the line.
+      borderColor: lineColor,
       tension: 0.2,
       pointRadius: Object.keys(data).indexOf(algorithm) === 0 ? 7 : 4,
       pointHoverRadius: Object.keys(data).indexOf(algorithm) === 0 ? 9 : 6,
       pointBackgroundColor: lineColor,
-      pointBorderColor: "#ffffff", // Add a border around points for contrast.
+      pointBorderColor: "#ffffff",
-      borderWidth: i === 0 ? 3 : 1.5, // Slightly increase line size for the primary dataset.
+      borderWidth: i === 0 ? 3 : 1.5,
-      hidden: false,
+      hidden:
+        initialDatasets.length > 0 && !initialDatasets.includes(algorithm),
     };
   });
 
@@ -152,7 +149,7 @@ function updateChart(initialDatasets = []) {
   modelComparisonChart = new Chart(
     document.getElementById("modelComparisonChart").getContext("2d"),
     {
-      type: "line", // Set the chart type to line.
+      type: "line",
       data: { datasets },
       options: {
         //aspectRatio: 2.5,  // higher is wider

tests/test_exports.py

@@ -44,18 +44,25 @@ def test_export_openvino():
 @pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
 @pytest.mark.parametrize(
     "task, dynamic, int8, half, batch, nms",
-    [  # generate all combinations but exclude those where both int8 and half are True
+    [  # generate all combinations except for exclusion cases
         (task, dynamic, int8, half, batch, nms)
         for task, dynamic, int8, half, batch, nms in product(
             TASKS, [True, False], [True, False], [True, False], [1, 2], [True, False]
         )
-        if not (int8 and half)  # exclude cases where both int8 and half are True
+        if not ((int8 and half) or (task == "classify" and nms))
     ],
 )
 def test_export_openvino_matrix(task, dynamic, int8, half, batch, nms):
     """Test YOLO model exports to OpenVINO under various configuration matrix conditions."""
     file = YOLO(TASK2MODEL[task]).export(
-        format="openvino", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, data=TASK2DATA[task], nms=nms
+        format="openvino",
+        imgsz=32,
+        dynamic=dynamic,
+        int8=int8,
+        half=half,
+        batch=batch,
+        data=TASK2DATA[task],
+        nms=nms,
     )
     if WINDOWS:
         # Use unique filenames due to Windows file permissions bug possibly due to latent threaded use
@@ -69,7 +76,13 @@ def test_export_openvino_matrix(task, dynamic, int8, half, batch, nms):
 @pytest.mark.slow
 @pytest.mark.parametrize(
     "task, dynamic, int8, half, batch, simplify, nms",
-    product(TASKS, [True, False], [False], [False], [1, 2], [True, False], [True, False]),
+    [  # generate all combinations except for exclusion cases
+        (task, dynamic, int8, half, batch, simplify, nms)
+        for task, dynamic, int8, half, batch, simplify, nms in product(
+            TASKS, [True, False], [False], [False], [1, 2], [True, False], [True, False]
+        )
+        if not ((int8 and half) or (task == "classify" and nms))
+    ],
 )
 def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
     """Test YOLO exports to ONNX format with various configurations and parameters."""
@@ -82,14 +95,19 @@ def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
 
 @pytest.mark.slow
 @pytest.mark.parametrize(
-    "task, dynamic, int8, half, batch, nms", product(TASKS, [False], [False], [False], [1, 2], [True, False])
+    "task, dynamic, int8, half, batch, nms",
+    [  # generate all combinations except for exclusion cases
+        (task, dynamic, int8, half, batch, nms)
+        for task, dynamic, int8, half, batch, nms in product(TASKS, [False], [False], [False], [1, 2], [True, False])
+        if not (task == "classify" and nms)
+    ],
 )
 def test_export_torchscript_matrix(task, dynamic, int8, half, batch, nms):
     """Tests YOLO model exports to TorchScript format under varied configurations."""
     file = YOLO(TASK2MODEL[task]).export(
         format="torchscript", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, nms=nms
     )
-    YOLO(file)([SOURCE] * 3, imgsz=64 if dynamic else 32)  # exported model inference at batch=3
+    YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32)  # exported model inference
     Path(file).unlink()  # cleanup
 
 
@@ -99,10 +117,10 @@ def test_export_torchscript_matrix(task, dynamic, int8, half, batch, nms):
 @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="CoreML not supported in Python 3.12")
 @pytest.mark.parametrize(
     "task, dynamic, int8, half, batch",
-    [  # generate all combinations but exclude those where both int8 and half are True
+    [  # generate all combinations except for exclusion cases
         (task, dynamic, int8, half, batch)
         for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
-        if not (int8 and half)  # exclude cases where both int8 and half are True
+        if not (int8 and half)
     ],
 )
 def test_export_coreml_matrix(task, dynamic, int8, half, batch):
@@ -124,12 +142,12 @@ def test_export_coreml_matrix(task, dynamic, int8, half, batch):
 @pytest.mark.skipif(not LINUX, reason="Test disabled as TF suffers from install conflicts on Windows and macOS")
 @pytest.mark.parametrize(
     "task, dynamic, int8, half, batch, nms",
-    [  # generate all combinations but exclude those where both int8 and half are True
+    [  # generate all combinations except for exclusion cases
         (task, dynamic, int8, half, batch, nms)
         for task, dynamic, int8, half, batch, nms in product(
             TASKS, [False], [True, False], [True, False], [1], [True, False]
         )
-        if not (int8 and half)  # exclude cases where both int8 and half are True
+        if not ((int8 and half) or (task == "classify" and nms))
     ],
 )
 def test_export_tflite_matrix(task, dynamic, int8, half, batch, nms):

ultralytics/engine/exporter.py

@@ -75,7 +75,7 @@ from ultralytics.data.dataset import YOLODataset
 from ultralytics.data.utils import check_cls_dataset, check_det_dataset
 from ultralytics.nn.autobackend import check_class_names, default_class_names
 from ultralytics.nn.modules import C2f, Classify, Detect, RTDETRDecoder
-from ultralytics.nn.tasks import DetectionModel, SegmentationModel, WorldModel
+from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, WorldModel
 from ultralytics.utils import (
     ARM64,
     DEFAULT_CFG,
@@ -282,6 +282,7 @@ class Exporter:
         if self.args.int8 and tflite:
             assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
         if self.args.nms:
+            assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
             if getattr(model, "end2end", False):
                 LOGGER.warning("WARNING ⚠️ 'nms=True' is not available for end2end models. Forcing 'nms=False'.")
                 self.args.nms = False
@@ -507,6 +508,7 @@ class Exporter:
         output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
         dynamic = self.args.dynamic
         if dynamic:
+            self.model.cpu()  # dynamic=True only compatible with cpu
             dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
             if isinstance(self.model, SegmentationModel):
                 dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 116, 8400)
@@ -518,13 +520,14 @@ class Exporter:
         if self.args.nms and self.model.task == "obb":
             self.args.opset = opset_version  # for NMSModel
             # OBB error https://github.com/pytorch/pytorch/issues/110859#issuecomment-1757841865
-            torch.onnx.register_custom_op_symbolic("aten::lift_fresh", lambda g, x: x, opset_version)
+            try:
+                torch.onnx.register_custom_op_symbolic("aten::lift_fresh", lambda g, x: x, opset_version)
+            except RuntimeError:  # it will fail if it's already registered
+                pass
             check_requirements("onnxslim>=0.1.46")  # Older versions has bug with OBB
 
         torch.onnx.export(
-            NMSModel(self.model.cpu() if dynamic else self.model, self.args)
+            NMSModel(self.model, self.args) if self.args.nms else self.model,
-            if self.args.nms
-            else self.model,  # dynamic=True only compatible with cpu
             self.im.cpu() if dynamic else self.im,
             f,
             verbose=False,
@@ -1570,7 +1573,7 @@ class NMSModel(torch.nn.Module):
                 # TFLite GatherND error if mask is empty
                 score *= mask
                 # Explicit length otherwise reshape error, hardcoded to `self.args.max_det * 5`
-                mask = score.topk(self.args.max_det * 5).indices
+                mask = score.topk(min(self.args.max_det * 5, score.shape[0])).indices
             box, score, cls, extra = box[mask], score[mask], cls[mask], extra[mask]
             if not self.obb:
                 box = xywh2xyxy(box)