⚡️ Speed up function output_to_target by 35%

[codeflash-ai]

📄 35% (0.35x) speedup for output_to_target in ultralytics/utils/plotting.py

⏱️ Runtime : 7.39 milliseconds → 5.48 milliseconds (best of 63 runs)

📝 Explanation and details

The optimization focuses on the xyxy2xywh function, which converts bounding box coordinates from (x1, y1, x2, y2) format to (x, y, width, height) format. The key improvement replaces four individual element-wise assignments with two vectorized slice operations.

What was optimized:

• Vectorized slice operations: Instead of assigning each coordinate individually (y[..., 0] = ..., y[..., 1] = ..., etc.), the optimized version uses slice assignments (y[..., :2] = ..., y[..., 2:] = ...) that operate on multiple elements simultaneously.
• Intermediate variable extraction: The coordinates are extracted once into xy and wh variables, reducing redundant indexing operations.

Why this leads to speedup:

• Reduced memory access: The original code performs 8 separate indexing operations (4 reads + 4 writes), while the optimized version performs 6 operations (4 reads + 2 writes).
• Better vectorization: PyTorch and NumPy are highly optimized for slice operations, which can leverage SIMD instructions and better memory access patterns compared to individual element assignments.
• Cache efficiency: Contiguous slice operations have better cache locality than scattered individual element access.

Performance impact in context:
The xyxy2xywh function is called from output_to_target, which is used in YOLO model validation for plotting predictions (as shown in the function references). During validation, this function processes detection results for visualization, and the 34% speedup directly reduces the time spent converting bounding box formats. The test results show consistent improvements across all scenarios, with particularly strong gains (39-51%) for large-scale cases with many batches or detections, making validation plotting significantly faster.

Test case benefits:
The optimization performs well across all test scenarios, with especially strong improvements for large-scale cases (many batches: 51% faster, large detections: 28-40% faster), indicating the vectorized approach scales better than individual assignments.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 31 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import torch
from ultralytics.utils.plotting import output_to_target

# Basic Test Cases


def test_single_batch_single_detection():
    # Test with a single batch and a single detection
    # Format: [x1, y1, x2, y2, conf, class]
    output = [torch.tensor([[10, 20, 30, 40, 0.9, 2]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 122μs -> 107μs (13.0% faster)
    x, y, w, h = boxes[0]


def test_multiple_batches_multiple_detections():
    # Test with two batches, each with two detections
    output = [
        torch.tensor([[0, 0, 10, 10, 0.8, 1], [10, 10, 20, 20, 0.7, 2]]),
        torch.tensor([[5, 5, 15, 15, 0.6, 3], [20, 20, 30, 30, 0.5, 4]]),
    ]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 149μs -> 123μs (21.0% faster)
    # Check box conversion
    expected_boxes = [
        [(0 + 10) / 2, (0 + 10) / 2, 10, 10],
        [(10 + 20) / 2, (10 + 20) / 2, 10, 10],
        [(5 + 15) / 2, (5 + 15) / 2, 10, 10],
        [(20 + 30) / 2, (20 + 30) / 2, 10, 10],
    ]
    for i, box in enumerate(boxes):
        pass


def test_max_det_limit():
    # Test that max_det limits the number of detections per batch
    output = [torch.tensor([[i, i + 1, i + 2, i + 3, 0.5, 1] for i in range(10)])]
    batch_id, class_id, boxes, conf = output_to_target(output, max_det=5)  # 105μs -> 89.8μs (17.1% faster)


# Edge Test Cases


def test_batch_with_no_detections():
    # Test with a batch that has no detections (empty tensor)
    output = [torch.empty((0, 6)), torch.tensor([[1, 2, 3, 4, 0.5, 1]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 158μs -> 134μs (18.0% faster)


def test_non_float_confidence_and_class():
    # Test with integer class and confidence as float
    output = [torch.tensor([[1, 2, 3, 4, 1.0, 0]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 95.0μs -> 84.0μs (13.0% faster)


def test_negative_coordinates():
    # Test with negative coordinates
    output = [torch.tensor([[-10, -20, 10, 20, 0.99, 1]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 96.1μs -> 78.5μs (22.5% faster)
    x, y, w, h = boxes[0]


def test_single_detection_max_det_greater_than_actual():
    # Test when max_det is greater than actual number of detections
    output = [torch.tensor([[1, 2, 3, 4, 0.5, 1]])]
    batch_id, class_id, boxes, conf = output_to_target(output, max_det=10)  # 93.1μs -> 80.3μs (15.9% faster)


def test_all_zero_confidence():
    # Test with all zero confidence scores
    output = [torch.tensor([[1, 2, 3, 4, 0.0, 1], [5, 6, 7, 8, 0.0, 2]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 99.6μs -> 83.8μs (18.8% faster)


def test_large_class_index():
    # Test with large class index
    output = [torch.tensor([[1, 2, 3, 4, 0.5, 999]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 92.0μs -> 81.5μs (12.9% faster)


def test_float_class_index():
    # Test with float class index
    output = [torch.tensor([[1, 2, 3, 4, 0.5, 2.5]])]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 94.2μs -> 80.9μs (16.5% faster)


def test_tensor_on_cuda():
    # Test with tensor on CUDA (if available)
    if torch.cuda.is_available():
        output = [torch.tensor([[1, 2, 3, 4, 0.5, 1]], device="cuda")]
        batch_id, class_id, boxes, conf = output_to_target(output)


# Large Scale Test Cases


def test_large_number_of_detections():
    # Test with a large number of detections in one batch
    num_det = 999
    output = [
        torch.cat(
            [
                torch.arange(num_det).unsqueeze(1).float(),  # x1
                torch.arange(num_det).unsqueeze(1).float() + 1,  # y1
                torch.arange(num_det).unsqueeze(1).float() + 2,  # x2
                torch.arange(num_det).unsqueeze(1).float() + 3,  # y2
                torch.ones((num_det, 1)),  # conf
                torch.zeros((num_det, 1)),  # class
            ],
            1,
        )
    ]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 93.3μs -> 80.8μs (15.4% faster)
    # Check that box conversion is correct for a few samples
    for i in [0, num_det // 2, num_det - 1]:
        x1, y1, x2, y2 = i, i + 1, i + 2, i + 3
        x = (x1 + x2) / 2
        y = (y1 + y2) / 2
        w = x2 - x1
        h = y2 - y1


def test_many_batches():
    # Test with many batches, each with a few detections
    num_batches = 10
    num_det = 10
    output = [
        torch.cat(
            [
                torch.full((num_det, 1), i),  # x1
                torch.full((num_det, 1), i + 1),  # y1
                torch.full((num_det, 1), i + 2),  # x2
                torch.full((num_det, 1), i + 3),  # y2
                torch.ones((num_det, 1)),  # conf
                torch.full((num_det, 1), i),  # class
            ],
            1,
        )
        for i in range(num_batches)
    ]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 440μs -> 315μs (39.7% faster)
    # Check batch indices
    for i in range(num_batches):
        start = i * num_det
        end = (i + 1) * num_det
        for j in range(start, end):
            x1 = i
            y1 = i + 1
            x2 = i + 2
            y2 = i + 3
            x = (x1 + x2) / 2
            y = (y1 + y2) / 2
            w = x2 - x1
            h = y2 - y1


def test_large_max_det():
    # Test with large max_det value
    output = [torch.tensor([[i, i + 1, i + 2, i + 3, 0.5, 1] for i in range(500)])]
    batch_id, class_id, boxes, conf = output_to_target(output, max_det=500)  # 117μs -> 106μs (11.1% faster)


def test_large_batch_and_large_detections():
    # Test with multiple batches, each with many detections
    num_batches = 5
    num_det = 200
    output = [
        torch.cat(
            [
                torch.arange(num_det).unsqueeze(1).float() + i * 1000,  # x1
                torch.arange(num_det).unsqueeze(1).float() + i * 1000 + 1,  # y1
                torch.arange(num_det).unsqueeze(1).float() + i * 1000 + 2,  # x2
                torch.arange(num_det).unsqueeze(1).float() + i * 1000 + 3,  # y2
                torch.ones((num_det, 1)),  # conf
                torch.full((num_det, 1), i),  # class
            ],
            1,
        )
        for i in range(num_batches)
    ]
    batch_id, class_id, boxes, conf = output_to_target(output)  # 249μs -> 194μs (28.3% faster)
    # Check batch indices and class indices for a few samples
    for i in range(num_batches):
        idx = i * num_det


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np

# imports
import pytest  # used for our unit tests
import torch
from ultralytics.utils.plotting import output_to_target

# ========== UNIT TESTS ==========

# -------- BASIC TEST CASES --------


def test_single_detection_single_batch():
    # Single batch, single detection, simple values
    # Format: [x1, y1, x2, y2, conf, class]
    det = torch.tensor([[10, 20, 30, 40, 0.9, 2]])
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 109μs -> 94.6μs (15.7% faster)


def test_multiple_detections_single_batch():
    # Single batch, multiple detections
    det = torch.tensor([[0, 0, 10, 10, 0.5, 1], [10, 10, 30, 30, 0.8, 3]])
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 107μs -> 86.2μs (24.5% faster)
    # Check xywh conversion
    expected_xywh = np.array([[5, 5, 10, 10], [20, 20, 20, 20]])


def test_multiple_batches():
    # Two batches, each with one detection
    det1 = torch.tensor([[1, 2, 3, 4, 0.7, 0]])
    det2 = torch.tensor([[5, 6, 7, 8, 0.8, 1]])
    output = [det1, det2]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 136μs -> 113μs (20.2% faster)


def test_max_det_limit():
    # Test that detections are limited to max_det
    det = torch.cat(
        [
            torch.arange(0, 600).reshape(100, 6).float(),  # 100 detections, dummy values
        ]
    )
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output, max_det=10)  # 87.2μs -> 75.5μs (15.5% faster)


# -------- EDGE TEST CASES --------


def test_empty_detections_in_batch():
    # One batch, zero detections
    det = torch.empty((0, 6))
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 115μs -> 102μs (13.0% faster)


def test_non_float_tensor_input():
    # Detections as integer tensor (should be cast to float32 internally)
    det = torch.tensor([[1, 2, 3, 4, 1, 0]], dtype=torch.int64)
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 106μs -> 88.7μs (19.9% faster)


def test_negative_and_zero_area_boxes():
    # Boxes with zero and negative area
    det = torch.tensor(
        [
            [1, 2, 1, 2, 0.5, 0],  # zero area
            [5, 6, 4, 5, 0.8, 1],  # negative area
        ]
    )
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 97.7μs -> 85.1μs (14.8% faster)


def test_high_class_index():
    # Class index is high (e.g., 999)
    det = torch.tensor([[0, 0, 1, 1, 0.99, 999]])
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 96.6μs -> 83.6μs (15.5% faster)


def test_input_on_cuda_if_available():
    # If CUDA is available, ensure function works with CUDA tensors
    if torch.cuda.is_available():
        det = torch.tensor([[1, 2, 3, 4, 0.5, 1]], device="cuda")
        output = [det]
        batch_id, class_id, xywh, conf = output_to_target(output)


def test_batch_with_varied_detection_counts():
    # Batches with different number of detections
    det1 = torch.tensor([[1, 2, 3, 4, 0.9, 2]])
    det2 = torch.tensor([[5, 6, 7, 8, 0.8, 1], [9, 10, 11, 12, 0.7, 3]])
    output = [det1, det2]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 170μs -> 142μs (19.2% faster)


def test_nan_and_inf_values():
    # Detections with NaN and Inf values
    det = torch.tensor(
        [[np.nan, 0, 1, 1, 0.5, 1], [0, np.inf, 1, 1, 0.6, 2], [0, 0, np.nan, 1, 0.7, 3], [0, 0, 1, np.inf, 0.8, 4]]
    )
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 100μs -> 84.7μs (19.2% faster)


# -------- LARGE SCALE TEST CASES --------


def test_large_number_of_detections():
    # Test with 1000 detections (limit to <100MB)
    n = 1000
    det = torch.zeros((n, 6), dtype=torch.float32)
    det[:, 0] = torch.arange(n)  # x1
    det[:, 1] = torch.arange(n)  # y1
    det[:, 2] = torch.arange(n) + 10  # x2
    det[:, 3] = torch.arange(n) + 20  # y2
    det[:, 4] = 0.5  # conf
    det[:, 5] = 1  # class
    output = [det]
    batch_id, class_id, xywh, conf = output_to_target(output)  # 89.6μs -> 78.9μs (13.5% faster)
    # Check a few values
    for idx in [0, n // 2, n - 1]:
        x1, y1, x2, y2 = det[idx, :4]
        expected_xy = [(x1 + x2) / 2, (y1 + y2) / 2, x2 - x1, y2 - y1]


def test_large_number_of_batches():
    # Test with 100 batches, each with 10 detections
    n_batches = 100
    n_det = 10
    output = []
    for i in range(n_batches):
        det = torch.zeros((n_det, 6), dtype=torch.float32)
        det[:, 0] = i
        det[:, 1] = i + 1
        det[:, 2] = i + 2
        det[:, 3] = i + 3
        det[:, 4] = 0.1 * i
        det[:, 5] = i % 5
        output.append(det)
    batch_id, class_id, xywh, conf = output_to_target(output)  # 3.98ms -> 2.64ms (51.0% faster)
    # Check batch id distribution
    for i in range(n_batches):
        idx_start = i * n_det
        idx_end = (i + 1) * n_det


def test_large_max_det_limit():
    # Test that max_det works for large input
    n = 500
    det = torch.zeros((n, 6), dtype=torch.float32)
    det[:, 0] = torch.arange(n)
    det[:, 1] = torch.arange(n)
    det[:, 2] = torch.arange(n) + 5
    det[:, 3] = torch.arange(n) + 10
    det[:, 4] = 0.9
    det[:, 5] = 7
    output = [det]
    # Only first 100 detections should be returned
    batch_id, class_id, xywh, conf = output_to_target(output, max_det=100)  # 88.8μs -> 76.3μs (16.4% faster)


# --------- FAILURE/MUTATION RESISTANCE TESTS ---------


def test_non_tensor_in_list():
    # Input list contains non-tensor
    output = [np.zeros((1, 6))]
    with pytest.raises(AttributeError):
        output_to_target(output)  # 3.54μs -> 3.52μs (0.483% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-output_to_target-mi8ewrk0 and push.