dgenerate-ultralytics-headless 8.3.197__py3-none-any.whl → 8.3.198__py3-none-any.whl

ultralytics/models/yolo/classify/val.py

@@ -178,7 +178,7 @@ class ClassificationValidator(BaseValidator):
             >>> batch = {"img": torch.rand(16, 3, 224, 224), "cls": torch.randint(0, 10, (16,))}
             >>> validator.plot_val_samples(batch, 0)
         """
-        batch["batch_idx"] = torch.arange(len(batch["img"]))  # add batch index for plotting
+        batch["batch_idx"] = torch.arange(batch["img"].shape[0])  # add batch index for plotting
         plot_images(
             labels=batch,
             fname=self.save_dir / f"val_batch{ni}_labels.jpg",
@@ -203,7 +203,7 @@ class ClassificationValidator(BaseValidator):
         """
         batched_preds = dict(
             img=batch["img"],
-            batch_idx=torch.arange(len(batch["img"])),
+            batch_idx=torch.arange(batch["img"].shape[0]),
             cls=torch.argmax(preds, dim=1),
         )
         plot_images(

ultralytics/models/yolo/detect/predict.py

@@ -89,7 +89,7 @@ class DetectionPredictor(BasePredictor):
         obj_feats = torch.cat(
             [x.permute(0, 2, 3, 1).reshape(x.shape[0], -1, s, x.shape[1] // s).mean(dim=-1) for x in feat_maps], dim=1
         )  # mean reduce all vectors to same length
-        return [feats[idx] if len(idx) else [] for feats, idx in zip(obj_feats, idxs)]  # for each img in batch
+        return [feats[idx] if idx.shape[0] else [] for feats, idx in zip(obj_feats, idxs)]  # for each img in batch
 
     def construct_results(self, preds, img, orig_imgs):
         """

ultralytics/models/yolo/detect/train.py

@@ -17,7 +17,7 @@ from ultralytics.models import yolo
 from ultralytics.nn.tasks import DetectionModel
 from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK
 from ultralytics.utils.patches import override_configs
-from ultralytics.utils.plotting import plot_images, plot_labels, plot_results
+from ultralytics.utils.plotting import plot_images, plot_labels
 from ultralytics.utils.torch_utils import torch_distributed_zero_first, unwrap_model
 
 
@@ -43,7 +43,6 @@ class DetectionTrainer(BaseTrainer):
         label_loss_items: Return a loss dictionary with labeled training loss items.
         progress_string: Return a formatted string of training progress.
         plot_training_samples: Plot training samples with their annotations.
-        plot_metrics: Plot metrics from a CSV file.
         plot_training_labels: Create a labeled training plot of the YOLO model.
         auto_batch: Calculate optimal batch size based on model memory requirements.
 
@@ -217,10 +216,6 @@ class DetectionTrainer(BaseTrainer):
             on_plot=self.on_plot,
         )
 
-    def plot_metrics(self):
-        """Plot metrics from a CSV file."""
-        plot_results(file=self.csv, on_plot=self.on_plot)  # save results.png
-
     def plot_training_labels(self):
         """Create a labeled training plot of the YOLO model."""
         boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0)

ultralytics/models/yolo/detect/val.py

@@ -146,7 +146,7 @@ class DetectionValidator(BaseValidator):
         ori_shape = batch["ori_shape"][si]
         imgsz = batch["img"].shape[2:]
         ratio_pad = batch["ratio_pad"][si]
-        if len(cls):
+        if cls.shape[0]:
             bbox = ops.xywh2xyxy(bbox) * torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]]  # target boxes
         return {
             "cls": cls,
@@ -185,7 +185,7 @@ class DetectionValidator(BaseValidator):
             predn = self._prepare_pred(pred)
 
             cls = pbatch["cls"].cpu().numpy()
-            no_pred = len(predn["cls"]) == 0
+            no_pred = predn["cls"].shape[0] == 0
             self.metrics.update_stats(
                 {
                     **self._process_batch(predn, pbatch),
@@ -268,8 +268,8 @@ class DetectionValidator(BaseValidator):
         Returns:
             (dict[str, np.ndarray]): Dictionary containing 'tp' key with correct prediction matrix of shape (N, 10) for 10 IoU levels.
         """
-        if len(batch["cls"]) == 0 or len(preds["cls"]) == 0:
-            return {"tp": np.zeros((len(preds["cls"]), self.niou), dtype=bool)}
+        if batch["cls"].shape[0] == 0 or preds["cls"].shape[0] == 0:
+            return {"tp": np.zeros((preds["cls"].shape[0], self.niou), dtype=bool)}
         iou = box_iou(batch["bboxes"], preds["bboxes"])
         return {"tp": self.match_predictions(preds["cls"], batch["cls"], iou).cpu().numpy()}
 

ultralytics/models/yolo/obb/val.py

@@ -93,8 +93,8 @@ class OBBValidator(DetectionValidator):
             >>> gt_cls = torch.randint(0, 5, (50,))  # 50 ground truth class labels
             >>> correct_matrix = validator._process_batch(detections, gt_bboxes, gt_cls)
         """
-        if len(batch["cls"]) == 0 or len(preds["cls"]) == 0:
-            return {"tp": np.zeros((len(preds["cls"]), self.niou), dtype=bool)}
+        if batch["cls"].shape[0] == 0 or preds["cls"].shape[0] == 0:
+            return {"tp": np.zeros((preds["cls"].shape[0], self.niou), dtype=bool)}
         iou = batch_probiou(batch["bboxes"], preds["bboxes"])
         return {"tp": self.match_predictions(preds["cls"], batch["cls"], iou).cpu().numpy()}
 
@@ -134,7 +134,7 @@ class OBBValidator(DetectionValidator):
         ori_shape = batch["ori_shape"][si]
         imgsz = batch["img"].shape[2:]
         ratio_pad = batch["ratio_pad"][si]
-        if len(cls):
+        if cls.shape[0]:
             bbox[..., :4].mul_(torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]])  # target boxes
         return {
             "cls": cls,

ultralytics/models/yolo/pose/predict.py

@@ -73,7 +73,7 @@ class PosePredictor(DetectionPredictor):
         """
         result = super().construct_result(pred, img, orig_img, img_path)
         # Extract keypoints from prediction and reshape according to model's keypoint shape
-        pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape)
+        pred_kpts = pred[:, 6:].view(pred.shape[0], *self.model.kpt_shape)
         # Scale keypoints coordinates to match the original image dimensions
         pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
         result.update(keypoints=pred_kpts)

ultralytics/models/yolo/pose/train.py

@@ -9,7 +9,6 @@ from typing import Any
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import PoseModel
 from ultralytics.utils import DEFAULT_CFG, LOGGER
-from ultralytics.utils.plotting import plot_results
 
 
 class PoseTrainer(yolo.detect.DetectionTrainer):
@@ -30,7 +29,6 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         set_model_attributes: Set keypoints shape attribute on the model.
         get_validator: Create a validator instance for model evaluation.
         plot_training_samples: Visualize training samples with keypoints.
-        plot_metrics: Generate and save training/validation metric plots.
         get_dataset: Retrieve the dataset and ensure it contains required kpt_shape key.
 
     Examples:
@@ -101,10 +99,6 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def plot_metrics(self):
-        """Plot training/validation metrics."""
-        plot_results(file=self.csv, pose=True, on_plot=self.on_plot)  # save results.png
-
     def get_dataset(self) -> dict[str, Any]:
         """
         Retrieve the dataset and ensure it contains the required `kpt_shape` key.

ultralytics/models/yolo/pose/val.py

@@ -192,8 +192,8 @@ class PoseValidator(DetectionValidator):
         """
         tp = super()._process_batch(preds, batch)
         gt_cls = batch["cls"]
-        if len(gt_cls) == 0 or len(preds["cls"]) == 0:
-            tp_p = np.zeros((len(preds["cls"]), self.niou), dtype=bool)
+        if gt_cls.shape[0] == 0 or preds["cls"].shape[0] == 0:
+            tp_p = np.zeros((preds["cls"].shape[0], self.niou), dtype=bool)
         else:
             # `0.53` is from https://github.com/jin-s13/xtcocoapi/blob/master/xtcocotools/cocoeval.py#L384
             area = ops.xyxy2xywh(batch["bboxes"])[:, 2:].prod(1) * 0.53

ultralytics/models/yolo/segment/predict.py

@@ -90,7 +90,7 @@ class SegmentationPredictor(DetectionPredictor):
         Construct a single result object from the prediction.
 
         Args:
-            pred (np.ndarray): The predicted bounding boxes, scores, and masks.
+            pred (torch.Tensor): The predicted bounding boxes, scores, and masks.
             img (torch.Tensor): The image after preprocessing.
             orig_img (np.ndarray): The original image before preprocessing.
             img_path (str): The path to the original image.
@@ -99,7 +99,7 @@ class SegmentationPredictor(DetectionPredictor):
         Returns:
             (Results): Result object containing the original image, image path, class names, bounding boxes, and masks.
         """
-        if not len(pred):  # save empty boxes
+        if pred.shape[0] == 0:  # save empty boxes
             masks = None
         elif self.args.retina_masks:
             pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)

ultralytics/models/yolo/segment/train.py

@@ -8,7 +8,6 @@ from pathlib import Path
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import SegmentationModel
 from ultralytics.utils import DEFAULT_CFG, RANK
-from ultralytics.utils.plotting import plot_results
 
 
 class SegmentationTrainer(yolo.detect.DetectionTrainer):
@@ -71,7 +70,3 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
         return yolo.segment.SegmentationValidator(
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
-
-    def plot_metrics(self):
-        """Plot training/validation metrics."""
-        plot_results(file=self.csv, segment=True, on_plot=self.on_plot)  # save results.png

ultralytics/models/yolo/segment/val.py

@@ -112,7 +112,7 @@ class SegmentationValidator(DetectionValidator):
             coefficient = pred.pop("extra")
             pred["masks"] = (
                 self.process(proto[i], coefficient, pred["bboxes"], shape=imgsz)
-                if len(coefficient)
+                if coefficient.shape[0]
                 else torch.zeros(
                     (0, *(imgsz if self.process is ops.process_mask_native else proto.shape[2:])),
                     dtype=torch.uint8,
@@ -133,16 +133,18 @@ class SegmentationValidator(DetectionValidator):
             (dict[str, Any]): Prepared batch with processed annotations.
         """
         prepared_batch = super()._prepare_batch(si, batch)
-        nl = len(prepared_batch["cls"])
+        nl = prepared_batch["cls"].shape[0]
         if self.args.overlap_mask:
             masks = batch["masks"][si]
             index = torch.arange(1, nl + 1, device=masks.device).view(nl, 1, 1)
             masks = (masks == index).float()
         else:
             masks = batch["masks"][batch["batch_idx"] == si]
-        if nl and self.process is ops.process_mask_native:
-            masks = F.interpolate(masks[None], prepared_batch["imgsz"], mode="bilinear", align_corners=False)[0]
-            masks = masks.gt_(0.5)
+        if nl:
+            mask_size = [s if self.process is ops.process_mask_native else s // 4 for s in prepared_batch["imgsz"]]
+            if masks.shape[1:] != mask_size:
+                masks = F.interpolate(masks[None], mask_size, mode="bilinear", align_corners=False)[0]
+                masks = masks.gt_(0.5)
         prepared_batch["masks"] = masks
         return prepared_batch
 
@@ -168,8 +170,8 @@ class SegmentationValidator(DetectionValidator):
         """
         tp = super()._process_batch(preds, batch)
         gt_cls = batch["cls"]
-        if len(gt_cls) == 0 or len(preds["cls"]) == 0:
-            tp_m = np.zeros((len(preds["cls"]), self.niou), dtype=bool)
+        if gt_cls.shape[0] == 0 or preds["cls"].shape[0] == 0:
+            tp_m = np.zeros((preds["cls"].shape[0], self.niou), dtype=bool)
         else:
             iou = mask_iou(batch["masks"].flatten(1), preds["masks"].flatten(1))
             tp_m = self.match_predictions(preds["cls"], gt_cls, iou).cpu().numpy()

ultralytics/models/yolo/yoloe/val.py

@@ -89,7 +89,7 @@ class YOLOEDetectValidator(DetectionValidator):
             for i in range(preds.shape[0]):
                 cls = batch["cls"][batch_idx == i].squeeze(-1).to(torch.int).unique(sorted=True)
                 pad_cls = torch.ones(preds.shape[1], device=self.device) * -1
-                pad_cls[: len(cls)] = cls
+                pad_cls[: cls.shape[0]] = cls
                 for c in cls:
                     visual_pe[c] += preds[i][pad_cls == c].sum(0) / cls_visual_num[c]
 

ultralytics/nn/modules/block.py

@@ -1921,7 +1921,7 @@ class A2C2f(nn.Module):
         y.extend(m(y[-1]) for m in self.m)
         y = self.cv2(torch.cat(y, 1))
         if self.gamma is not None:
-            return x + self.gamma.view(-1, len(self.gamma), 1, 1) * y
+            return x + self.gamma.view(-1, self.gamma.shape[0], 1, 1) * y
         return y
 
 

ultralytics/nn/tasks.py

@@ -766,7 +766,7 @@ class RTDETRDetectionModel(DetectionModel):
 
         img = batch["img"]
         # NOTE: preprocess gt_bbox and gt_labels to list.
-        bs = len(img)
+        bs = img.shape[0]
         batch_idx = batch["batch_idx"]
         gt_groups = [(batch_idx == i).sum().item() for i in range(bs)]
         targets = {
@@ -923,7 +923,7 @@ class WorldModel(DetectionModel):
             (torch.Tensor): Model's output tensor.
         """
         txt_feats = (self.txt_feats if txt_feats is None else txt_feats).to(device=x.device, dtype=x.dtype)
-        if len(txt_feats) != len(x) or self.model[-1].export:
+        if txt_feats.shape[0] != x.shape[0] or self.model[-1].export:
             txt_feats = txt_feats.expand(x.shape[0], -1, -1)
         ori_txt_feats = txt_feats.clone()
         y, dt, embeddings = [], [], []  # outputs

ultralytics/utils/checks.py

@@ -907,7 +907,7 @@ def is_intel():
     try:
         result = subprocess.run(["xpu-smi", "discovery"], capture_output=True, text=True, timeout=5)
         return "intel" in result.stdout.lower()
-    except (subprocess.TimeoutExpired, FileNotFoundError, subprocess.SubprocessError):
+    except Exception:  # broad clause to capture all Intel GPU exception types
         return False
 
 

ultralytics/utils/metrics.py

@@ -397,11 +397,11 @@ class ConfusionMatrix(DataExportMixin):
         gt_cls, gt_bboxes = batch["cls"], batch["bboxes"]
         if self.matches is not None:  # only if visualization is enabled
             self.matches = {k: defaultdict(list) for k in {"TP", "FP", "FN", "GT"}}
-            for i in range(len(gt_cls)):
+            for i in range(gt_cls.shape[0]):
                 self._append_matches("GT", batch, i)  # store GT
         is_obb = gt_bboxes.shape[1] == 5  # check if boxes contains angle for OBB
         conf = 0.25 if conf in {None, 0.01 if is_obb else 0.001} else conf  # apply 0.25 if default val conf is passed
-        no_pred = len(detections["cls"]) == 0
+        no_pred = detections["cls"].shape[0] == 0
         if gt_cls.shape[0] == 0:  # Check if labels is empty
             if not no_pred:
                 detections = {k: detections[k][detections["conf"] > conf] for k in detections}
@@ -491,13 +491,13 @@ class ConfusionMatrix(DataExportMixin):
         for i, mtype in enumerate(["GT", "FP", "TP", "FN"]):
             mbatch = self.matches[mtype]
             if "conf" not in mbatch:
-                mbatch["conf"] = torch.tensor([1.0] * len(mbatch["bboxes"]), device=img.device)
-            mbatch["batch_idx"] = torch.ones(len(mbatch["bboxes"]), device=img.device) * i
+                mbatch["conf"] = torch.tensor([1.0] * mbatch["bboxes"].shape[0], device=img.device)
+            mbatch["batch_idx"] = torch.ones(mbatch["bboxes"].shape[0], device=img.device) * i
             for k in mbatch.keys():
                 labels[k] += mbatch[k]
 
         labels = {k: torch.stack(v, 0) if len(v) else v for k, v in labels.items()}
-        if self.task != "obb" and len(labels["bboxes"]):
+        if self.task != "obb" and labels["bboxes"].shape[0]:
             labels["bboxes"] = xyxy2xywh(labels["bboxes"])
         (save_dir / "visualizations").mkdir(parents=True, exist_ok=True)
         plot_images(
@@ -980,7 +980,7 @@ class Metric(SimpleClass):
 
     def fitness(self) -> float:
         """Return model fitness as a weighted combination of metrics."""
-        w = [0.0, 0.0, 0.1, 0.9]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
+        w = [0.0, 0.0, 0.0, 1.0]  # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
         return (np.nan_to_num(np.array(self.mean_results())) * w).sum()
 
     def update(self, results: tuple):

ultralytics/utils/nms.py

@@ -263,12 +263,11 @@ class TorchNMS:
         areas = (x2 - x1) * (y2 - y1)
 
         # Sort by scores descending
-        _, order = scores.sort(0, descending=True)
+        order = scores.argsort(0, descending=True)
 
         # Pre-allocate keep list with maximum possible size
         keep = torch.zeros(order.numel(), dtype=torch.int64, device=boxes.device)
         keep_idx = 0
-
         while order.numel() > 0:
             i = order[0]
             keep[keep_idx] = i
@@ -276,7 +275,6 @@ class TorchNMS:
 
             if order.numel() == 1:
                 break
-
             # Vectorized IoU calculation for remaining boxes
             rest = order[1:]
             xx1 = torch.maximum(x1[i], x1[rest])
@@ -288,20 +286,14 @@ class TorchNMS:
             w = (xx2 - xx1).clamp_(min=0)
             h = (yy2 - yy1).clamp_(min=0)
             inter = w * h
-
-            # Early termination: skip IoU calculation if no intersection
+            # Early exit: skip IoU calculation if no intersection
             if inter.sum() == 0:
                 # No overlaps with current box, keep all remaining boxes
-                remaining_count = rest.numel()
-                keep[keep_idx : keep_idx + remaining_count] = rest
-                keep_idx += remaining_count
-                break
-
+                order = rest
+                continue
             iou = inter / (areas[i] + areas[rest] - inter)
-
             # Keep boxes with IoU <= threshold
-            mask = iou <= iou_threshold
-            order = rest[mask]
+            order = rest[iou <= iou_threshold]
 
         return keep[:keep_idx]
 

ultralytics/utils/plotting.py

@@ -812,14 +812,13 @@ def plot_images(
 
             # Plot masks
             if len(masks):
-                if idx.shape[0] == masks.shape[0]:  # overlap_mask=False
+                if idx.shape[0] == masks.shape[0] and masks.max() <= 1:  # overlap_mask=False
                     image_masks = masks[idx]
                 else:  # overlap_mask=True
                     image_masks = masks[[i]]  # (1, 640, 640)
                     nl = idx.sum()
-                    index = np.arange(nl).reshape((nl, 1, 1)) + 1
-                    image_masks = np.repeat(image_masks, nl, axis=0)
-                    image_masks = np.where(image_masks == index, 1.0, 0.0)
+                    index = np.arange(1, nl + 1).reshape((nl, 1, 1))
+                    image_masks = (image_masks == index).astype(np.float32)
 
                 im = np.asarray(annotator.im).copy()
                 for j in range(len(image_masks)):
@@ -847,14 +846,7 @@ def plot_images(
 
 
 @plt_settings()
-def plot_results(
-    file: str = "path/to/results.csv",
-    dir: str = "",
-    segment: bool = False,
-    pose: bool = False,
-    classify: bool = False,
-    on_plot: Callable | None = None,
-):
+def plot_results(file: str = "path/to/results.csv", dir: str = "", on_plot: Callable | None = None):
     """
     Plot training results from a results CSV file. The function supports various types of data including segmentation,
     pose estimation, and classification. Plots are saved as 'results.png' in the directory where the CSV is located.
@@ -862,9 +854,6 @@ def plot_results(
     Args:
         file (str, optional): Path to the CSV file containing the training results.
         dir (str, optional): Directory where the CSV file is located if 'file' is not provided.
-        segment (bool, optional): Flag to indicate if the data is for segmentation.
-        pose (bool, optional): Flag to indicate if the data is for pose estimation.
-        classify (bool, optional): Flag to indicate if the data is for classification.
         on_plot (callable, optional): Callback function to be executed after plotting. Takes filename as an argument.
 
     Examples:
@@ -876,34 +865,31 @@ def plot_results(
     from scipy.ndimage import gaussian_filter1d
 
     save_dir = Path(file).parent if file else Path(dir)
-    if classify:
-        fig, ax = plt.subplots(2, 2, figsize=(6, 6), tight_layout=True)
-        index = [2, 5, 3, 4]
-    elif segment:
-        fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
-        index = [2, 3, 4, 5, 6, 7, 10, 11, 14, 15, 16, 17, 8, 9, 12, 13]
-    elif pose:
-        fig, ax = plt.subplots(2, 9, figsize=(21, 6), tight_layout=True)
-        index = [2, 3, 4, 5, 6, 7, 8, 11, 12, 15, 16, 17, 18, 19, 9, 10, 13, 14]
-    else:
-        fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
-        index = [2, 3, 4, 5, 6, 9, 10, 11, 7, 8]
-    ax = ax.ravel()
     files = list(save_dir.glob("results*.csv"))
     assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
-    for f in files:
+
+    loss_keys, metric_keys = [], []
+    for i, f in enumerate(files):
         try:
             data = pl.read_csv(f, infer_schema_length=None)
-            s = [x.strip() for x in data.columns]
+            if i == 0:
+                for c in data.columns:
+                    if "loss" in c:
+                        loss_keys.append(c)
+                    elif "metric" in c:
+                        metric_keys.append(c)
+                loss_mid, metric_mid = len(loss_keys) // 2, len(metric_keys) // 2
+                columns = (
+                    loss_keys[:loss_mid] + metric_keys[:metric_mid] + loss_keys[loss_mid:] + metric_keys[metric_mid:]
+                )
+                fig, ax = plt.subplots(2, len(columns) // 2, figsize=(len(columns) + 2, 6), tight_layout=True)
+                ax = ax.ravel()
             x = data.select(data.columns[0]).to_numpy().flatten()
-            for i, j in enumerate(index):
-                y = data.select(data.columns[j]).to_numpy().flatten().astype("float")
-                # y[y == 0] = np.nan  # don't show zero values
+            for i, j in enumerate(columns):
+                y = data.select(j).to_numpy().flatten().astype("float")
                 ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8)  # actual results
                 ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2)  # smoothing line
-                ax[i].set_title(s[j], fontsize=12)
-                # if j in {8, 9, 10}:  # share train and val loss y axes
-                #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
+                ax[i].set_title(j, fontsize=12)
         except Exception as e:
             LOGGER.error(f"Plotting error for {f}: {e}")
     ax[1].legend()

ultralytics/utils/torch_utils.py

@@ -1012,7 +1012,7 @@ def attempt_compile(
     imgsz: int = 640,
     use_autocast: bool = False,
     warmup: bool = False,
-    prefix: str = colorstr("compile:"),
+    mode: bool | str = "default",
 ) -> torch.nn.Module:
     """
     Compile a model with torch.compile and optionally warm up the graph to reduce first-iteration latency.
@@ -1027,7 +1027,8 @@ def attempt_compile(
         imgsz (int, optional): Square input size to create a dummy tensor with shape (1, 3, imgsz, imgsz) for warmup.
         use_autocast (bool, optional): Whether to run warmup under autocast on CUDA or MPS devices.
         warmup (bool, optional): Whether to execute a single dummy forward pass to warm up the compiled model.
-        prefix (str, optional): Message prefix for logger output.
+        mode (bool | str, optional): torch.compile mode. True → "default", False → no compile, or a string like
+            "default", "reduce-overhead", "max-autotune".
 
     Returns:
         model (torch.nn.Module): Compiled model if compilation succeeds, otherwise the original unmodified model.
@@ -1042,13 +1043,16 @@ def attempt_compile(
         >>> # Try to compile and warm up a model with a 640x640 input
         >>> model = attempt_compile(model, device=device, imgsz=640, use_autocast=True, warmup=True)
     """
-    if not hasattr(torch, "compile"):
+    if not hasattr(torch, "compile") or not mode:
         return model
 
-    LOGGER.info(f"{prefix} starting torch.compile...")
+    if mode is True:
+        mode = "default"
+    prefix = colorstr("compile:")
+    LOGGER.info(f"{prefix} starting torch.compile with '{mode}' mode...")
     t0 = time.perf_counter()
     try:
-        model = torch.compile(model, mode="max-autotune", backend="inductor")
+        model = torch.compile(model, mode=mode, backend="inductor")
     except Exception as e:
         LOGGER.warning(f"{prefix} torch.compile failed, continuing uncompiled: {e}")
         return model