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>
This commit is contained in:
parent
83dc1fea6e
commit
de05d1b655
3 changed files with 46 additions and 28 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue