ultralytics 8.3.30__py3-none-any.whl → 8.3.31__py3-none-any.whl

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.3.30"
+__version__ = "8.3.31"
 
 import os
 

ultralytics/engine/trainer.py

@@ -279,12 +279,7 @@ class BaseTrainer:
 
         # Batch size
         if self.batch_size < 1 and RANK == -1:  # single-GPU only, estimate best batch size
-            self.args.batch = self.batch_size = check_train_batch_size(
-                model=self.model,
-                imgsz=self.args.imgsz,
-                amp=self.amp,
-                batch=self.batch_size,
-            )
+            self.args.batch = self.batch_size = self.auto_batch()
 
         # Dataloaders
         batch_size = self.batch_size // max(world_size, 1)
@@ -478,6 +473,16 @@ class BaseTrainer:
         self._clear_memory()
         self.run_callbacks("teardown")
 
+    def auto_batch(self, max_num_obj=0):
+        """Get batch size by calculating memory occupation of model."""
+        return check_train_batch_size(
+            model=self.model,
+            imgsz=self.args.imgsz,
+            amp=self.amp,
+            batch=self.batch_size,
+            max_num_obj=max_num_obj,
+        )  # returns batch size
+
     def _get_memory(self):
         """Get accelerator memory utilization in GB."""
         if self.device.type == "mps":

ultralytics/models/yolo/detect/train.py

@@ -141,3 +141,10 @@ class DetectionTrainer(BaseTrainer):
         boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0)
         cls = np.concatenate([lb["cls"] for lb in self.train_loader.dataset.labels], 0)
         plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)
+
+    def auto_batch(self):
+        """Get batch size by calculating memory occupation of model."""
+        train_dataset = self.build_dataset(self.trainset, mode="train", batch=16)
+        # 4 for mosaic augmentation
+        max_num_obj = max(len(l["cls"]) for l in train_dataset.labels) * 4
+        return super().auto_batch(max_num_obj)

ultralytics/utils/autobatch.py

@@ -11,7 +11,7 @@ from ultralytics.utils import DEFAULT_CFG, LOGGER, colorstr
 from ultralytics.utils.torch_utils import autocast, profile
 
 
-def check_train_batch_size(model, imgsz=640, amp=True, batch=-1):
+def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):
     """
     Compute optimal YOLO training batch size using the autobatch() function.
 
@@ -20,6 +20,7 @@ def check_train_batch_size(model, imgsz=640, amp=True, batch=-1):
         imgsz (int, optional): Image size used for training.
         amp (bool, optional): Use automatic mixed precision if True.
         batch (float, optional): Fraction of GPU memory to use. If -1, use default.
+        max_num_obj (int, optional): The maximum number of objects from dataset.
 
     Returns:
         (int): Optimal batch size computed using the autobatch() function.
@@ -29,10 +30,12 @@ def check_train_batch_size(model, imgsz=640, amp=True, batch=-1):
         Otherwise, a default fraction of 0.6 is used.
     """
     with autocast(enabled=amp):
-        return autobatch(deepcopy(model).train(), imgsz, fraction=batch if 0.0 < batch < 1.0 else 0.6)
+        return autobatch(
+            deepcopy(model).train(), imgsz, fraction=batch if 0.0 < batch < 1.0 else 0.6, max_num_obj=max_num_obj
+        )
 
 
-def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
+def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch, max_num_obj=1):
     """
     Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
 
@@ -41,6 +44,7 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
         imgsz (int, optional): The image size used as input for the YOLO model. Defaults to 640.
         fraction (float, optional): The fraction of available CUDA memory to use. Defaults to 0.60.
         batch_size (int, optional): The default batch size to use if an error is detected. Defaults to 16.
+        max_num_obj (int, optional): The maximum number of objects from dataset.
 
     Returns:
         (int): The optimal batch size.
@@ -70,7 +74,7 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
     batch_sizes = [1, 2, 4, 8, 16] if t < 16 else [1, 2, 4, 8, 16, 32, 64]
     try:
         img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
-        results = profile(img, model, n=1, device=device)
+        results = profile(img, model, n=1, device=device, max_num_obj=max_num_obj)
 
         # Fit a solution
         y = [x[2] for x in results if x]  # memory [2]

ultralytics/utils/tal.py

@@ -3,6 +3,7 @@
 import torch
 import torch.nn as nn
 
+from . import LOGGER
 from .checks import check_version
 from .metrics import bbox_iou, probiou
 from .ops import xywhr2xyxyxyxy
@@ -73,6 +74,7 @@ class TaskAlignedAssigner(nn.Module):
             return self._forward(pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt)
         except torch.OutOfMemoryError:
             # Move tensors to CPU, compute, then move back to original device
+            LOGGER.warning("WARNING: CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU")
             cpu_tensors = [t.cpu() for t in (pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt)]
             result = self._forward(*cpu_tensors)
             return tuple(t.to(device) for t in result)

ultralytics/utils/torch_utils.py

@@ -623,7 +623,7 @@ def convert_optimizer_state_dict_to_fp16(state_dict):
     return state_dict
 
 
-def profile(input, ops, n=10, device=None):
+def profile(input, ops, n=10, device=None, max_num_obj=0):
     """
     Ultralytics speed, memory and FLOPs profiler.
 
@@ -671,6 +671,14 @@ def profile(input, ops, n=10, device=None):
                         t[2] = float("nan")
                     tf += (t[1] - t[0]) * 1000 / n  # ms per op forward
                     tb += (t[2] - t[1]) * 1000 / n  # ms per op backward
+                    if max_num_obj:  # simulate training with predictions per image grid (for AutoBatch)
+                        torch.randn(
+                            x.shape[0],
+                            max_num_obj,
+                            int(sum([(x.shape[-1] / s) * (x.shape[-2] / s) for s in m.stride.tolist()])),
+                            device=device,
+                            dtype=torch.float32,
+                        )
                 mem = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0  # (GB)
                 s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else "list" for x in (x, y))  # shapes
                 p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0  # parameters

{ultralytics-8.3.30.dist-info → ultralytics-8.3.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics
-Version: 8.3.30
+Version: 8.3.31
 Summary: Ultralytics YOLO 🚀 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{ultralytics-8.3.30.dist-info → ultralytics-8.3.31.dist-info}/RECORD

@@ -7,7 +7,7 @@ tests/test_exports.py,sha256=1MvhcQ2qHdbJImHII-bFarcaIcm-kPlEK-OdFLxnj7o,8769
 tests/test_integrations.py,sha256=f5-QCUk1SU_-qn4mBCZwS3GN3tXEBIIXo4z2EhExbHw,6126
 tests/test_python.py,sha256=I1RRdCwLdrc3jX06huVxct8HX8ccQOmQgVpuEflRl0U,23560
 tests/test_solutions.py,sha256=sPYhy2d814mIVvojQeVxeZPu0IVy01_Y8zuMcu_9GF0,3790
-ultralytics/__init__.py,sha256=rWIiTJOQfPzIOhdWwz2WFyiM9SD4GYuZ2eYzbj6MQYQ,681
+ultralytics/__init__.py,sha256=Z3VjR0LGu2YpNK8RouNCLcCrmcAaslU5opcp7PtMzsg,681
 ultralytics/assets/bus.jpg,sha256=wCAZxJecGR63Od3ZRERe9Aja1Weayrb9Ug751DS_vGM,137419
 ultralytics/assets/zidane.jpg,sha256=Ftc4aeMmen1O0A3o6GCDO9FlfBslLpTAw0gnetx7bts,50427
 ultralytics/cfg/__init__.py,sha256=ArJow4-pOPN3y6aKOv5KcVXimikI6vAQvQlSRb7IdWE,38743
@@ -103,7 +103,7 @@ ultralytics/engine/exporter.py,sha256=DH67LwNDr3fiWxaES-lhSLvm5pCuasXLbQv4FSLCi_
 ultralytics/engine/model.py,sha256=TfuTczFjNJ3GW0E_qWVH6OaJ_2I-_Srx7i_4GQebDoo,51472
 ultralytics/engine/predictor.py,sha256=aS4yJdTK2kYq-TTpzIlWxqnAcBz38zIECZoMb_yOPMY,17597
 ultralytics/engine/results.py,sha256=a1XFZRPwqgKDBOEAibHuT9nP2xefLiWVsMoBJbcr4iA,75058
-ultralytics/engine/trainer.py,sha256=lbFMLdrdWkk1td6BpUS0_uLhAkiWo-eAmx_Kaov1JPA,37149
+ultralytics/engine/trainer.py,sha256=Cd95QLJ3C4fncoOX1YgauLA9aWVYRd1G6x0Au2xX86k,37335
 ultralytics/engine/tuner.py,sha256=WBj8iw1K1TK0hvanlA-wkwmfqh1SI8jEe2dGwUINeTg,11838
 ultralytics/engine/validator.py,sha256=aWpXE3nrOqaA7jCuUgwxi0FabiGTIXtZvjoJyCX903o,14870
 ultralytics/hub/__init__.py,sha256=c6Me4E8V-P7mtzTggyPYz9FnVkqWRyPp9F-fMcyFNQ0,5632
@@ -151,7 +151,7 @@ ultralytics/models/yolo/classify/train.py,sha256=3aYzLDqX_03xR1xqlTn1TxA4t58cCIG
 ultralytics/models/yolo/classify/val.py,sha256=Tzizhp3ebzPvwJejrE8tb-TuXw4MdkEI9mOANV74eXQ,4909
 ultralytics/models/yolo/detect/__init__.py,sha256=JR8gZJWn7wMBbh-0j_073nxJVZTMFZVWTOG5Wnvk6w0,229
 ultralytics/models/yolo/detect/predict.py,sha256=-uZFLutxGYZX47RANcaxC-LFStRbv0nBv_8-ypadQoI,1471
-ultralytics/models/yolo/detect/train.py,sha256=jayPByykRhJvUeDCN-oZ4M6Yx0rpjAWq9epjJgnHeYA,6363
+ultralytics/models/yolo/detect/train.py,sha256=f-QwAxaRsdtY-KY9femcWLQt1wB1tAHu-QV0a13LOGA,6702
 ultralytics/models/yolo/detect/val.py,sha256=Ydf0W7FPf-nSz_lOCTcIrlDqhG_kOK8Od6eYsGQDx9U,15144
 ultralytics/models/yolo/obb/__init__.py,sha256=txWbPGLY1_M7ZwlLQjrwGjTBOlsv9P3yk5ZEgysTinU,193
 ultralytics/models/yolo/obb/predict.py,sha256=VxpKCKV5dWnOr0GyV1rJGH5SzzRouCYW_8T26xJ8MU8,2037
@@ -200,7 +200,7 @@ ultralytics/trackers/utils/gmc.py,sha256=VcURuY041qGCeWUGMxHZBr10T16LtcMqyv7AmTf
 ultralytics/trackers/utils/kalman_filter.py,sha256=cH9zD3fwkuezP97H9mw8cSBN7a8hHKx_Sx1j7t3oYGs,21349
 ultralytics/trackers/utils/matching.py,sha256=3Ie1WNNRZ4_q3365F03XD7Nr9juZB_08mw4yUKC3w74,7162
 ultralytics/utils/__init__.py,sha256=08pFkzKn1eR9xdIFhx8tx_8MO-gqXjt2n0HGwDeUlWE,49159
-ultralytics/utils/autobatch.py,sha256=BO9MCRtrLDtrDQaxqV0BdjaYsgXf-q07Y3_VdGp4URY,4330
+ultralytics/utils/autobatch.py,sha256=nt0nSNNhrQqvtaxeNBBYpU2OkZnI3ihNEAa3jF4pybo,4594
 ultralytics/utils/benchmarks.py,sha256=aEW28iVIMj-8bwOgISDphOJExDmaGi5bz3G2PJlRjcc,25793
 ultralytics/utils/checks.py,sha256=KXQSeauhzecy9tSjyDVy8oXbTDkHSSB9lOTYrqRWpok,29582
 ultralytics/utils/dist.py,sha256=NDFga-uKxkBX2zLxFHSene_cCiGQJoyOeCXcN9JIOIk,2358
@@ -213,8 +213,8 @@ ultralytics/utils/metrics.py,sha256=msPaXc244ndc0NPBhnNlHsKkVhdc-TMgFn5NATlZZVI,
 ultralytics/utils/ops.py,sha256=dsXNdyrYx_p6io6zezig9p84dxS7U-10vceHNVu2IL0,32888
 ultralytics/utils/patches.py,sha256=J-iOwIRbfUs-inBZerhnXby5tUKjYcOIyvhLTS352JE,3270
 ultralytics/utils/plotting.py,sha256=TKtdbAOl6gZdFD2hlA5T4LNWfr2LUWbCC-cXkgL1JAU,61089
-ultralytics/utils/tal.py,sha256=t5sTZwaoglx3pW0P5ootiKFL-AsKZksh3ypHbPTInKE,18354
-ultralytics/utils/torch_utils.py,sha256=jB03Q-9ajTplxE05CdkmJmpXDUkb4LSiv3S6S2laWII,31608
+ultralytics/utils/tal.py,sha256=thD_AEhVmhaZqmS5szZMvpKO-RKOeZwfX1BYAhdnA0o,18470
+ultralytics/utils/torch_utils.py,sha256=CfftLVZqCEgIOLtRsU4mytRJoHDiXd5ijzIqIVt-zZM,32070
 ultralytics/utils/triton.py,sha256=gg1finxno_tY2Ge9PMhmu7PI9wvoFZoiicdT4Bhqv3w,3936
 ultralytics/utils/tuner.py,sha256=K09-z5k1E4ZriSKoWdwQrJ2PJ2fY1ez3-b2R6aKPTqM,6198
 ultralytics/utils/callbacks/__init__.py,sha256=YrWqC3BVVaTLob4iCPR6I36mUxIUOpPJW7B_LjT78Qw,214
@@ -228,9 +228,9 @@ ultralytics/utils/callbacks/neptune.py,sha256=IbGQfEltamUKXJt93uSLQFn8c2rYh3DMTg
 ultralytics/utils/callbacks/raytune.py,sha256=Ck_yFzg7UZXiDWrLHaltjQybzVWSFDfzpdrx9ZYTRfI,700
 ultralytics/utils/callbacks/tensorboard.py,sha256=SHlE58Fb-sg-uZKtgy-ybIO3SAIfK55aj8kTYGA0Cyg,4167
 ultralytics/utils/callbacks/wb.py,sha256=oX3JarCJGhzvW556XiEXQNaZblAaK_UETAt3kzkY61w,6869
-ultralytics-8.3.30.dist-info/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
-ultralytics-8.3.30.dist-info/METADATA,sha256=Sfqm2741qh7purvJHD8A_h1wQs4Gb72aQVH1H1eXg4A,35213
-ultralytics-8.3.30.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-ultralytics-8.3.30.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
-ultralytics-8.3.30.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
-ultralytics-8.3.30.dist-info/RECORD,,
+ultralytics-8.3.31.dist-info/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
+ultralytics-8.3.31.dist-info/METADATA,sha256=pCRq3cWOvmofljRT9bsJ0P_x0Mtk2uNQXQUHZf8scO8,35213
+ultralytics-8.3.31.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+ultralytics-8.3.31.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
+ultralytics-8.3.31.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
+ultralytics-8.3.31.dist-info/RECORD,,