ultralytics 8.0.237__py3-none-any.whl → 8.0.239__py3-none-any.whl

ultralytics/utils/tal.py

@@ -7,7 +7,7 @@ from .checks import check_version
 from .metrics import bbox_iou, probiou
 from .ops import xywhr2xyxyxyxy
 
-TORCH_1_10 = check_version(torch.__version__, '1.10.0')
+TORCH_1_10 = check_version(torch.__version__, "1.10.0")
 
 
 class TaskAlignedAssigner(nn.Module):
@@ -61,12 +61,17 @@ class TaskAlignedAssigner(nn.Module):
 
         if self.n_max_boxes == 0:
             device = gt_bboxes.device
-            return (torch.full_like(pd_scores[..., 0], self.bg_idx).to(device), torch.zeros_like(pd_bboxes).to(device),
-                    torch.zeros_like(pd_scores).to(device), torch.zeros_like(pd_scores[..., 0]).to(device),
-                    torch.zeros_like(pd_scores[..., 0]).to(device))
-
-        mask_pos, align_metric, overlaps = self.get_pos_mask(pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points,
-                                                             mask_gt)
+            return (
+                torch.full_like(pd_scores[..., 0], self.bg_idx).to(device),
+                torch.zeros_like(pd_bboxes).to(device),
+                torch.zeros_like(pd_scores).to(device),
+                torch.zeros_like(pd_scores[..., 0]).to(device),
+                torch.zeros_like(pd_scores[..., 0]).to(device),
+            )
+
+        mask_pos, align_metric, overlaps = self.get_pos_mask(
+            pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt
+        )
 
         target_gt_idx, fg_mask, mask_pos = self.select_highest_overlaps(mask_pos, overlaps, self.n_max_boxes)
 
@@ -148,7 +153,7 @@ class TaskAlignedAssigner(nn.Module):
         ones = torch.ones_like(topk_idxs[:, :, :1], dtype=torch.int8, device=topk_idxs.device)
         for k in range(self.topk):
             # Expand topk_idxs for each value of k and add 1 at the specified positions
-            count_tensor.scatter_add_(-1, topk_idxs[:, :, k:k + 1], ones)
+            count_tensor.scatter_add_(-1, topk_idxs[:, :, k : k + 1], ones)
         # count_tensor.scatter_add_(-1, topk_idxs, torch.ones_like(topk_idxs, dtype=torch.int8, device=topk_idxs.device))
         # Filter invalid bboxes
         count_tensor.masked_fill_(count_tensor > 1, 0)
@@ -192,9 +197,11 @@ class TaskAlignedAssigner(nn.Module):
         target_labels.clamp_(0)
 
         # 10x faster than F.one_hot()
-        target_scores = torch.zeros((target_labels.shape[0], target_labels.shape[1], self.num_classes),
-                                    dtype=torch.int64,
-                                    device=target_labels.device)  # (b, h*w, 80)
+        target_scores = torch.zeros(
+            (target_labels.shape[0], target_labels.shape[1], self.num_classes),
+            dtype=torch.int64,
+            device=target_labels.device,
+        )  # (b, h*w, 80)
         target_scores.scatter_(2, target_labels.unsqueeze(-1), 1)
 
         fg_scores_mask = fg_mask[:, :, None].repeat(1, 1, self.num_classes)  # (b, h*w, 80)
@@ -252,7 +259,6 @@ class TaskAlignedAssigner(nn.Module):
 
 
 class RotatedTaskAlignedAssigner(TaskAlignedAssigner):
-
     def iou_calculation(self, gt_bboxes, pd_bboxes):
         """Iou calculation for rotated bounding boxes."""
         return probiou(gt_bboxes, pd_bboxes).squeeze(-1).clamp_(0)
@@ -295,7 +301,7 @@ def make_anchors(feats, strides, grid_cell_offset=0.5):
         _, _, h, w = feats[i].shape
         sx = torch.arange(end=w, device=device, dtype=dtype) + grid_cell_offset  # shift x
         sy = torch.arange(end=h, device=device, dtype=dtype) + grid_cell_offset  # shift y
-        sy, sx = torch.meshgrid(sy, sx, indexing='ij') if TORCH_1_10 else torch.meshgrid(sy, sx)
+        sy, sx = torch.meshgrid(sy, sx, indexing="ij") if TORCH_1_10 else torch.meshgrid(sy, sx)
         anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2))
         stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))
     return torch.cat(anchor_points), torch.cat(stride_tensor)

ultralytics/utils/torch_utils.py

@@ -25,11 +25,11 @@ try:
 except ImportError:
     thop = None
 
-TORCH_1_9 = check_version(torch.__version__, '1.9.0')
-TORCH_2_0 = check_version(torch.__version__, '2.0.0')
-TORCHVISION_0_10 = check_version(torchvision.__version__, '0.10.0')
-TORCHVISION_0_11 = check_version(torchvision.__version__, '0.11.0')
-TORCHVISION_0_13 = check_version(torchvision.__version__, '0.13.0')
+TORCH_1_9 = check_version(torch.__version__, "1.9.0")
+TORCH_2_0 = check_version(torch.__version__, "2.0.0")
+TORCHVISION_0_10 = check_version(torchvision.__version__, "0.10.0")
+TORCHVISION_0_11 = check_version(torchvision.__version__, "0.11.0")
+TORCHVISION_0_13 = check_version(torchvision.__version__, "0.13.0")
 
 
 @contextmanager
@@ -60,13 +60,13 @@ def get_cpu_info():
     """Return a string with system CPU information, i.e. 'Apple M2'."""
     import cpuinfo  # pip install py-cpuinfo
 
-    k = 'brand_raw', 'hardware_raw', 'arch_string_raw'  # info keys sorted by preference (not all keys always available)
+    k = "brand_raw", "hardware_raw", "arch_string_raw"  # info keys sorted by preference (not all keys always available)
     info = cpuinfo.get_cpu_info()  # info dict
-    string = info.get(k[0] if k[0] in info else k[1] if k[1] in info else k[2], 'unknown')
-    return string.replace('(R)', '').replace('CPU ', '').replace('@ ', '')
+    string = info.get(k[0] if k[0] in info else k[1] if k[1] in info else k[2], "unknown")
+    return string.replace("(R)", "").replace("CPU ", "").replace("@ ", "")
 
 
-def select_device(device='', batch=0, newline=False, verbose=True):
+def select_device(device="", batch=0, newline=False, verbose=True):
     """
     Selects the appropriate PyTorch device based on the provided arguments.
 
@@ -103,49 +103,57 @@ def select_device(device='', batch=0, newline=False, verbose=True):
     if isinstance(device, torch.device):
         return device
 
-    s = f'Ultralytics YOLOv{__version__} 🚀 Python-{platform.python_version()} torch-{torch.__version__} '
+    s = f"Ultralytics YOLOv{__version__} 🚀 Python-{platform.python_version()} torch-{torch.__version__} "
     device = str(device).lower()
-    for remove in 'cuda:', 'none', '(', ')', '[', ']', "'", ' ':
-        device = device.replace(remove, '')  # to string, 'cuda:0' -> '0' and '(0, 1)' -> '0,1'
-    cpu = device == 'cpu'
-    mps = device in ('mps', 'mps:0')  # Apple Metal Performance Shaders (MPS)
+    for remove in "cuda:", "none", "(", ")", "[", "]", "'", " ":
+        device = device.replace(remove, "")  # to string, 'cuda:0' -> '0' and '(0, 1)' -> '0,1'
+    cpu = device == "cpu"
+    mps = device in ("mps", "mps:0")  # Apple Metal Performance Shaders (MPS)
     if cpu or mps:
-        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # force torch.cuda.is_available() = False
     elif device:  # non-cpu device requested
-        if device == 'cuda':
-            device = '0'
-        visible = os.environ.get('CUDA_VISIBLE_DEVICES', None)
-        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable - must be before assert is_available()
-        if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', ''))):
+        if device == "cuda":
+            device = "0"
+        visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+        os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable - must be before assert is_available()
+        if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(",", ""))):
             LOGGER.info(s)
-            install = 'See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no ' \
-                      'CUDA devices are seen by torch.\n' if torch.cuda.device_count() == 0 else ''
-            raise ValueError(f"Invalid CUDA 'device={device}' requested."
-                             f" Use 'device=cpu' or pass valid CUDA device(s) if available,"
-                             f" i.e. 'device=0' or 'device=0,1,2,3' for Multi-GPU.\n"
-                             f'\ntorch.cuda.is_available(): {torch.cuda.is_available()}'
-                             f'\ntorch.cuda.device_count(): {torch.cuda.device_count()}'
-                             f"\nos.environ['CUDA_VISIBLE_DEVICES']: {visible}\n"
-                             f'{install}')
+            install = (
+                "See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no "
+                "CUDA devices are seen by torch.\n"
+                if torch.cuda.device_count() == 0
+                else ""
+            )
+            raise ValueError(
+                f"Invalid CUDA 'device={device}' requested."
+                f" Use 'device=cpu' or pass valid CUDA device(s) if available,"
+                f" i.e. 'device=0' or 'device=0,1,2,3' for Multi-GPU.\n"
+                f"\ntorch.cuda.is_available(): {torch.cuda.is_available()}"
+                f"\ntorch.cuda.device_count(): {torch.cuda.device_count()}"
+                f"\nos.environ['CUDA_VISIBLE_DEVICES']: {visible}\n"
+                f"{install}"
+            )
 
     if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available
-        devices = device.split(',') if device else '0'  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
+        devices = device.split(",") if device else "0"  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
         n = len(devices)  # device count
         if n > 1 and batch > 0 and batch % n != 0:  # check batch_size is divisible by device_count
-            raise ValueError(f"'batch={batch}' must be a multiple of GPU count {n}. Try 'batch={batch // n * n}' or "
-                             f"'batch={batch // n * n + n}', the nearest batch sizes evenly divisible by {n}.")
-        space = ' ' * (len(s) + 1)
+            raise ValueError(
+                f"'batch={batch}' must be a multiple of GPU count {n}. Try 'batch={batch // n * n}' or "
+                f"'batch={batch // n * n + n}', the nearest batch sizes evenly divisible by {n}."
+            )
+        space = " " * (len(s) + 1)
         for i, d in enumerate(devices):
             p = torch.cuda.get_device_properties(i)
             s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n"  # bytes to MB
-        arg = 'cuda:0'
+        arg = "cuda:0"
     elif mps and TORCH_2_0 and torch.backends.mps.is_available():
         # Prefer MPS if available
-        s += f'MPS ({get_cpu_info()})\n'
-        arg = 'mps'
+        s += f"MPS ({get_cpu_info()})\n"
+        arg = "mps"
     else:  # revert to CPU
-        s += f'CPU ({get_cpu_info()})\n'
-        arg = 'cpu'
+        s += f"CPU ({get_cpu_info()})\n"
+        arg = "cpu"
 
     if verbose:
         LOGGER.info(s if newline else s.rstrip())
@@ -161,14 +169,20 @@ def time_sync():
 
 def fuse_conv_and_bn(conv, bn):
     """Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/."""
-    fusedconv = nn.Conv2d(conv.in_channels,
-                          conv.out_channels,
-                          kernel_size=conv.kernel_size,
-                          stride=conv.stride,
-                          padding=conv.padding,
-                          dilation=conv.dilation,
-                          groups=conv.groups,
-                          bias=True).requires_grad_(False).to(conv.weight.device)
+    fusedconv = (
+        nn.Conv2d(
+            conv.in_channels,
+            conv.out_channels,
+            kernel_size=conv.kernel_size,
+            stride=conv.stride,
+            padding=conv.padding,
+            dilation=conv.dilation,
+            groups=conv.groups,
+            bias=True,
+        )
+        .requires_grad_(False)
+        .to(conv.weight.device)
+    )
 
     # Prepare filters
     w_conv = conv.weight.clone().view(conv.out_channels, -1)
@@ -185,15 +199,21 @@ def fuse_conv_and_bn(conv, bn):
 
 def fuse_deconv_and_bn(deconv, bn):
     """Fuse ConvTranspose2d() and BatchNorm2d() layers."""
-    fuseddconv = nn.ConvTranspose2d(deconv.in_channels,
-                                    deconv.out_channels,
-                                    kernel_size=deconv.kernel_size,
-                                    stride=deconv.stride,
-                                    padding=deconv.padding,
-                                    output_padding=deconv.output_padding,
-                                    dilation=deconv.dilation,
-                                    groups=deconv.groups,
-                                    bias=True).requires_grad_(False).to(deconv.weight.device)
+    fuseddconv = (
+        nn.ConvTranspose2d(
+            deconv.in_channels,
+            deconv.out_channels,
+            kernel_size=deconv.kernel_size,
+            stride=deconv.stride,
+            padding=deconv.padding,
+            output_padding=deconv.output_padding,
+            dilation=deconv.dilation,
+            groups=deconv.groups,
+            bias=True,
+        )
+        .requires_grad_(False)
+        .to(deconv.weight.device)
+    )
 
     # Prepare filters
     w_deconv = deconv.weight.clone().view(deconv.out_channels, -1)
@@ -221,18 +241,21 @@ def model_info(model, detailed=False, verbose=True, imgsz=640):
     n_l = len(list(model.modules()))  # number of layers
     if detailed:
         LOGGER.info(
-            f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}")
+            f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}"
+        )
         for i, (name, p) in enumerate(model.named_parameters()):
-            name = name.replace('module_list.', '')
-            LOGGER.info('%5g %40s %9s %12g %20s %10.3g %10.3g %10s' %
-                        (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std(), p.dtype))
+            name = name.replace("module_list.", "")
+            LOGGER.info(
+                "%5g %40s %9s %12g %20s %10.3g %10.3g %10s"
+                % (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std(), p.dtype)
+            )
 
     flops = get_flops(model, imgsz)
-    fused = ' (fused)' if getattr(model, 'is_fused', lambda: False)() else ''
-    fs = f', {flops:.1f} GFLOPs' if flops else ''
-    yaml_file = getattr(model, 'yaml_file', '') or getattr(model, 'yaml', {}).get('yaml_file', '')
-    model_name = Path(yaml_file).stem.replace('yolo', 'YOLO') or 'Model'
-    LOGGER.info(f'{model_name} summary{fused}: {n_l} layers, {n_p} parameters, {n_g} gradients{fs}')
+    fused = " (fused)" if getattr(model, "is_fused", lambda: False)() else ""
+    fs = f", {flops:.1f} GFLOPs" if flops else ""
+    yaml_file = getattr(model, "yaml_file", "") or getattr(model, "yaml", {}).get("yaml_file", "")
+    model_name = Path(yaml_file).stem.replace("yolo", "YOLO") or "Model"
+    LOGGER.info(f"{model_name} summary{fused}: {n_l} layers, {n_p} parameters, {n_g} gradients{fs}")
     return n_l, n_p, n_g, flops
 
 
@@ -262,13 +285,15 @@ def model_info_for_loggers(trainer):
     """
     if trainer.args.profile:  # profile ONNX and TensorRT times
         from ultralytics.utils.benchmarks import ProfileModels
+
         results = ProfileModels([trainer.last], device=trainer.device).profile()[0]
-        results.pop('model/name')
+        results.pop("model/name")
     else:  # only return PyTorch times from most recent validation
         results = {
-            'model/parameters': get_num_params(trainer.model),
-            'model/GFLOPs': round(get_flops(trainer.model), 3)}
-    results['model/speed_PyTorch(ms)'] = round(trainer.validator.speed['inference'], 3)
+            "model/parameters": get_num_params(trainer.model),
+            "model/GFLOPs": round(get_flops(trainer.model), 3),
+        }
+    results["model/speed_PyTorch(ms)"] = round(trainer.validator.speed["inference"], 3)
     return results
 
 
@@ -284,14 +309,14 @@ def get_flops(model, imgsz=640):
             imgsz = [imgsz, imgsz]  # expand if int/float
         try:
             # Use stride size for input tensor
-            stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32  # max stride
+            stride = max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32  # max stride
             im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
-            flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1E9 * 2  # stride GFLOPs
+            flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2  # stride GFLOPs
             return flops * imgsz[0] / stride * imgsz[1] / stride  # imgsz GFLOPs
         except Exception:
             # Use actual image size for input tensor (i.e. required for RTDETR models)
             im = torch.empty((1, p.shape[1], *imgsz), device=p.device)  # input image in BCHW format
-            return thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1E9 * 2  # imgsz GFLOPs
+            return thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2  # imgsz GFLOPs
     except Exception:
         return 0.0
 
@@ -301,11 +326,11 @@ def get_flops_with_torch_profiler(model, imgsz=640):
     if TORCH_2_0:
         model = de_parallel(model)
         p = next(model.parameters())
-        stride = (max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32) * 2  # max stride
+        stride = (max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32) * 2  # max stride
         im = torch.zeros((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
         with torch.profiler.profile(with_flops=True) as prof:
             model(im)
-        flops = sum(x.flops for x in prof.key_averages()) / 1E9
+        flops = sum(x.flops for x in prof.key_averages()) / 1e9
         imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz]  # expand if int/float
         flops = flops * imgsz[0] / stride * imgsz[1] / stride  # 640x640 GFLOPs
         return flops
@@ -333,7 +358,7 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32):
         return img
     h, w = img.shape[2:]
     s = (int(h * ratio), int(w * ratio))  # new size
-    img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
+    img = F.interpolate(img, size=s, mode="bilinear", align_corners=False)  # resize
     if not same_shape:  # pad/crop img
         h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
     return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
@@ -349,7 +374,7 @@ def make_divisible(x, divisor):
 def copy_attr(a, b, include=(), exclude=()):
     """Copies attributes from object 'b' to object 'a', with options to include/exclude certain attributes."""
     for k, v in b.__dict__.items():
-        if (len(include) and k not in include) or k.startswith('_') or k in exclude:
+        if (len(include) and k not in include) or k.startswith("_") or k in exclude:
             continue
         else:
             setattr(a, k, v)
@@ -357,7 +382,7 @@ def copy_attr(a, b, include=(), exclude=()):
 
 def get_latest_opset():
     """Return second-most (for maturity) recently supported ONNX opset by this version of torch."""
-    return max(int(k[14:]) for k in vars(torch.onnx) if 'symbolic_opset' in k) - 1  # opset
+    return max(int(k[14:]) for k in vars(torch.onnx) if "symbolic_opset" in k) - 1  # opset
 
 
 def intersect_dicts(da, db, exclude=()):
@@ -392,10 +417,10 @@ def init_seeds(seed=0, deterministic=False):
         if TORCH_2_0:
             torch.use_deterministic_algorithms(True, warn_only=True)  # warn if deterministic is not possible
             torch.backends.cudnn.deterministic = True
-            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
-            os.environ['PYTHONHASHSEED'] = str(seed)
+            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+            os.environ["PYTHONHASHSEED"] = str(seed)
         else:
-            LOGGER.warning('WARNING ⚠️ Upgrade to torch>=2.0.0 for deterministic training.')
+            LOGGER.warning("WARNING ⚠️ Upgrade to torch>=2.0.0 for deterministic training.")
     else:
         torch.use_deterministic_algorithms(False)
         torch.backends.cudnn.deterministic = False
@@ -430,13 +455,13 @@ class ModelEMA:
                     v += (1 - d) * msd[k].detach()
                     # assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype},  model {msd[k].dtype}'
 
-    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
+    def update_attr(self, model, include=(), exclude=("process_group", "reducer")):
         """Updates attributes and saves stripped model with optimizer removed."""
         if self.enabled:
             copy_attr(self.ema, model, include, exclude)
 
 
-def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None:
+def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "") -> None:
     """
     Strip optimizer from 'f' to finalize training, optionally save as 's'.
 
@@ -456,26 +481,26 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None:
             strip_optimizer(f)
         ```
     """
-    x = torch.load(f, map_location=torch.device('cpu'))
-    if 'model' not in x:
-        LOGGER.info(f'Skipping {f}, not a valid Ultralytics model.')
+    x = torch.load(f, map_location=torch.device("cpu"))
+    if "model" not in x:
+        LOGGER.info(f"Skipping {f}, not a valid Ultralytics model.")
         return
 
-    if hasattr(x['model'], 'args'):
-        x['model'].args = dict(x['model'].args)  # convert from IterableSimpleNamespace to dict
-    args = {**DEFAULT_CFG_DICT, **x['train_args']} if 'train_args' in x else None  # combine args
-    if x.get('ema'):
-        x['model'] = x['ema']  # replace model with ema
-    for k in 'optimizer', 'best_fitness', 'ema', 'updates':  # keys
+    if hasattr(x["model"], "args"):
+        x["model"].args = dict(x["model"].args)  # convert from IterableSimpleNamespace to dict
+    args = {**DEFAULT_CFG_DICT, **x["train_args"]} if "train_args" in x else None  # combine args
+    if x.get("ema"):
+        x["model"] = x["ema"]  # replace model with ema
+    for k in "optimizer", "best_fitness", "ema", "updates":  # keys
         x[k] = None
-    x['epoch'] = -1
-    x['model'].half()  # to FP16
-    for p in x['model'].parameters():
+    x["epoch"] = -1
+    x["model"].half()  # to FP16
+    for p in x["model"].parameters():
         p.requires_grad = False
-    x['train_args'] = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS}  # strip non-default keys
+    x["train_args"] = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS}  # strip non-default keys
     # x['model'].args = x['train_args']
     torch.save(x, s or f)
-    mb = os.path.getsize(s or f) / 1E6  # file size
+    mb = os.path.getsize(s or f) / 1e6  # file size
     LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
 
 
@@ -496,18 +521,20 @@ def profile(input, ops, n=10, device=None):
     results = []
     if not isinstance(device, torch.device):
         device = select_device(device)
-    LOGGER.info(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
-                f"{'input':>24s}{'output':>24s}")
+    LOGGER.info(
+        f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
+        f"{'input':>24s}{'output':>24s}"
+    )
 
     for x in input if isinstance(input, list) else [input]:
         x = x.to(device)
         x.requires_grad = True
         for m in ops if isinstance(ops, list) else [ops]:
-            m = m.to(device) if hasattr(m, 'to') else m  # device
-            m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
+            m = m.to(device) if hasattr(m, "to") else m  # device
+            m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
             tf, tb, t = 0, 0, [0, 0, 0]  # dt forward, backward
             try:
-                flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1E9 * 2 if thop else 0  # GFLOPs
+                flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1e9 * 2 if thop else 0  # GFLOPs
             except Exception:
                 flops = 0
 
@@ -521,13 +548,13 @@ def profile(input, ops, n=10, device=None):
                         t[2] = time_sync()
                     except Exception:  # no backward method
                         # print(e)  # for debug
-                        t[2] = float('nan')
+                        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
-                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
+                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
-                LOGGER.info(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}')
+                LOGGER.info(f"{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}")
                 results.append([p, flops, mem, tf, tb, s_in, s_out])
             except Exception as e:
                 LOGGER.info(e)
@@ -548,7 +575,7 @@ class EarlyStopping:
         """
         self.best_fitness = 0.0  # i.e. mAP
         self.best_epoch = 0
-        self.patience = patience or float('inf')  # epochs to wait after fitness stops improving to stop
+        self.patience = patience or float("inf")  # epochs to wait after fitness stops improving to stop
         self.possible_stop = False  # possible stop may occur next epoch
 
     def __call__(self, epoch, fitness):
@@ -572,8 +599,10 @@ class EarlyStopping:
         self.possible_stop = delta >= (self.patience - 1)  # possible stop may occur next epoch
         stop = delta >= self.patience  # stop training if patience exceeded
         if stop:
-            LOGGER.info(f'Stopping training early as no improvement observed in last {self.patience} epochs. '
-                        f'Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n'
-                        f'To update EarlyStopping(patience={self.patience}) pass a new patience value, '
-                        f'i.e. `patience=300` or use `patience=0` to disable EarlyStopping.')
+            LOGGER.info(
+                f"Stopping training early as no improvement observed in last {self.patience} epochs. "
+                f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n"
+                f"To update EarlyStopping(patience={self.patience}) pass a new patience value, "
+                f"i.e. `patience=300` or use `patience=0` to disable EarlyStopping."
+            )
         return stop

ultralytics/utils/triton.py

@@ -22,7 +22,7 @@ class TritonRemoteModel:
         output_names (List[str]): The names of the model outputs.
     """
 
-    def __init__(self, url: str, endpoint: str = '', scheme: str = ''):
+    def __init__(self, url: str, endpoint: str = "", scheme: str = ""):
         """
         Initialize the TritonRemoteModel.
 
@@ -36,7 +36,7 @@ class TritonRemoteModel:
         """
         if not endpoint and not scheme:  # Parse all args from URL string
             splits = urlsplit(url)
-            endpoint = splits.path.strip('/').split('/')[0]
+            endpoint = splits.path.strip("/").split("/")[0]
             scheme = splits.scheme
             url = splits.netloc
 
@@ -44,26 +44,28 @@ class TritonRemoteModel:
         self.url = url
 
         # Choose the Triton client based on the communication scheme
-        if scheme == 'http':
+        if scheme == "http":
             import tritonclient.http as client  # noqa
+
             self.triton_client = client.InferenceServerClient(url=self.url, verbose=False, ssl=False)
             config = self.triton_client.get_model_config(endpoint)
         else:
             import tritonclient.grpc as client  # noqa
+
             self.triton_client = client.InferenceServerClient(url=self.url, verbose=False, ssl=False)
-            config = self.triton_client.get_model_config(endpoint, as_json=True)['config']
+            config = self.triton_client.get_model_config(endpoint, as_json=True)["config"]
 
         # Sort output names alphabetically, i.e. 'output0', 'output1', etc.
-        config['output'] = sorted(config['output'], key=lambda x: x.get('name'))
+        config["output"] = sorted(config["output"], key=lambda x: x.get("name"))
 
         # Define model attributes
-        type_map = {'TYPE_FP32': np.float32, 'TYPE_FP16': np.float16, 'TYPE_UINT8': np.uint8}
+        type_map = {"TYPE_FP32": np.float32, "TYPE_FP16": np.float16, "TYPE_UINT8": np.uint8}
         self.InferRequestedOutput = client.InferRequestedOutput
         self.InferInput = client.InferInput
-        self.input_formats = [x['data_type'] for x in config['input']]
+        self.input_formats = [x["data_type"] for x in config["input"]]
         self.np_input_formats = [type_map[x] for x in self.input_formats]
-        self.input_names = [x['name'] for x in config['input']]
-        self.output_names = [x['name'] for x in config['output']]
+        self.input_names = [x["name"] for x in config["input"]]
+        self.output_names = [x["name"] for x in config["output"]]
 
     def __call__(self, *inputs: np.ndarray) -> List[np.ndarray]:
         """
@@ -80,7 +82,7 @@ class TritonRemoteModel:
         for i, x in enumerate(inputs):
             if x.dtype != self.np_input_formats[i]:
                 x = x.astype(self.np_input_formats[i])
-            infer_input = self.InferInput(self.input_names[i], [*x.shape], self.input_formats[i].replace('TYPE_', ''))
+            infer_input = self.InferInput(self.input_names[i], [*x.shape], self.input_formats[i].replace("TYPE_", ""))
             infer_input.set_data_from_numpy(x)
             infer_inputs.append(infer_input)