From de05d1b65509e14ff9396a3aa5ba8a958bcedcf9 Mon Sep 17 00:00:00 2001 From: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Date: Sun, 26 Jan 2025 04:26:58 +0800 Subject: [PATCH] Fix export test matrices to exclude `nms` from Classify models (#18880) Signed-off-by: Glenn Jocher Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- docs/overrides/javascript/benchmark.js | 21 ++++++-------- tests/test_exports.py | 38 +++++++++++++++++++------- ultralytics/engine/exporter.py | 15 ++++++---- 3 files changed, 46 insertions(+), 28 deletions(-) diff --git a/docs/overrides/javascript/benchmark.js b/docs/overrides/javascript/benchmark.js index f81e27ec..024514d3 100644 --- a/docs/overrides/javascript/benchmark.js +++ b/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. - const selectedAlgorithms = - initialDatasets.length > 0 ? initialDatasets : Object.keys(data); - - // Create the datasets for the selected algorithms. - const datasets = selectedAlgorithms.map((algorithm, i) => { + // Always include all models in the dataset creation + const datasets = Object.keys(data).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. - borderWidth: i === 0 ? 3 : 1.5, // Slightly increase line size for the primary dataset. - hidden: false, + pointBorderColor: "#ffffff", + borderWidth: i === 0 ? 3 : 1.5, + 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 diff --git a/tests/test_exports.py b/tests/test_exports.py index c364b481..3c96f7fe 100644 --- a/tests/test_exports.py +++ b/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): diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 88e6e533..90807106 100644 --- a/ultralytics/engine/exporter.py +++ b/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) - if self.args.nms - else self.model, # dynamic=True only compatible with cpu + NMSModel(self.model, self.args) if self.args.nms else self.model, 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)