dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.4.7__py3-none-any.whl

ultralytics/nn/text_model.py

@@ -20,8 +20,7 @@ except ImportError:
 
 
 class TextModel(nn.Module):
-    """
-    Abstract base class for text encoding models.
+    """Abstract base class for text encoding models.
 
     This class defines the interface for text encoding models used in vision-language tasks. Subclasses must implement
     the tokenize and encode_text methods to provide text tokenization and encoding functionality.
@@ -47,11 +46,10 @@ class TextModel(nn.Module):
 
 
 class CLIP(TextModel):
-    """
-    Implements OpenAI's CLIP (Contrastive Language-Image Pre-training) text encoder.
+    """Implements OpenAI's CLIP (Contrastive Language-Image Pre-training) text encoder.
 
-    This class provides a text encoder based on OpenAI's CLIP model, which can convert text into feature vectors
-    that are aligned with corresponding image features in a shared embedding space.
+    This class provides a text encoder based on OpenAI's CLIP model, which can convert text into feature vectors that
+    are aligned with corresponding image features in a shared embedding space.
 
     Attributes:
         model (clip.model.CLIP): The loaded CLIP model.
@@ -71,20 +69,14 @@ class CLIP(TextModel):
     """
 
     def __init__(self, size: str, device: torch.device) -> None:
-        """
-        Initialize the CLIP text encoder.
+        """Initialize the CLIP text encoder.
 
-        This class implements the TextModel interface using OpenAI's CLIP model for text encoding. It loads
-        a pre-trained CLIP model of the specified size and prepares it for text encoding tasks.
+        This class implements the TextModel interface using OpenAI's CLIP model for text encoding. It loads a
+        pre-trained CLIP model of the specified size and prepares it for text encoding tasks.
 
         Args:
             size (str): Model size identifier (e.g., 'ViT-B/32').
             device (torch.device): Device to load the model on.
-
-        Examples:
-            >>> import torch
-            >>> clip_model = CLIP("ViT-B/32", device=torch.device("cuda:0"))
-            >>> text_features = clip_model.encode_text(["a photo of a cat", "a photo of a dog"])
         """
         super().__init__()
         self.model, self.image_preprocess = clip.load(size, device=device)
@@ -92,12 +84,13 @@ class CLIP(TextModel):
         self.device = device
         self.eval()
 
-    def tokenize(self, texts: str | list[str]) -> torch.Tensor:
-        """
-        Convert input texts to CLIP tokens.
+    def tokenize(self, texts: str | list[str], truncate: bool = True) -> torch.Tensor:
+        """Convert input texts to CLIP tokens.
 
         Args:
             texts (str | list[str]): Input text or list of texts to tokenize.
+            truncate (bool, optional): Whether to trim texts that exceed CLIP's context length. Defaults to True to
+                avoid RuntimeError from overly long inputs while still allowing explicit opt-out.
 
         Returns:
             (torch.Tensor): Tokenized text tensor with shape (batch_size, context_length) ready for model processing.
@@ -106,13 +99,14 @@ class CLIP(TextModel):
             >>> model = CLIP("ViT-B/32", device="cpu")
             >>> tokens = model.tokenize("a photo of a cat")
             >>> print(tokens.shape)  # torch.Size([1, 77])
+            >>> strict_tokens = model.tokenize("a photo of a cat", truncate=False)  # Enforce strict length checks
+            >>> print(strict_tokens.shape)  # Same shape/content as tokens since prompt less than 77 tokens
         """
-        return clip.tokenize(texts).to(self.device)
+        return clip.tokenize(texts, truncate=truncate).to(self.device)
 
     @smart_inference_mode()
     def encode_text(self, texts: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
-        """
-        Encode tokenized texts into normalized feature vectors.
+        """Encode tokenized texts into normalized feature vectors.
 
         This method processes tokenized text inputs through the CLIP model to generate feature vectors, which are then
         normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.
@@ -137,15 +131,14 @@ class CLIP(TextModel):
 
     @smart_inference_mode()
     def encode_image(self, image: Image.Image | torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
-        """
-        Encode preprocessed images into normalized feature vectors.
+        """Encode preprocessed images into normalized feature vectors.
 
-        This method processes preprocessed image inputs through the CLIP model to generate feature vectors, which are then
-        normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.
+        This method processes preprocessed image inputs through the CLIP model to generate feature vectors, which are
+        then normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.
 
         Args:
-            image (PIL.Image | torch.Tensor): Preprocessed image input. If a PIL Image is provided, it will be
-                converted to a tensor using the model's image preprocessing function.
+            image (PIL.Image | torch.Tensor): Preprocessed image input. If a PIL Image is provided, it will be converted
+                to a tensor using the model's image preprocessing function.
             dtype (torch.dtype, optional): Data type for output features.
 
         Returns:
@@ -169,8 +162,7 @@ class CLIP(TextModel):
 
 
 class MobileCLIP(TextModel):
-    """
-    Implement Apple's MobileCLIP text encoder for efficient text encoding.
+    """Implement Apple's MobileCLIP text encoder for efficient text encoding.
 
     This class implements the TextModel interface using Apple's MobileCLIP model, providing efficient text encoding
     capabilities for vision-language tasks with reduced computational requirements compared to standard CLIP models.
@@ -195,28 +187,16 @@ class MobileCLIP(TextModel):
     config_size_map = {"s0": "s0", "s1": "s1", "s2": "s2", "b": "b", "blt": "b"}
 
     def __init__(self, size: str, device: torch.device) -> None:
-        """
-        Initialize the MobileCLIP text encoder.
+        """Initialize the MobileCLIP text encoder.
 
         This class implements the TextModel interface using Apple's MobileCLIP model for efficient text encoding.
 
         Args:
             size (str): Model size identifier (e.g., 's0', 's1', 's2', 'b', 'blt').
             device (torch.device): Device to load the model on.
-
-        Examples:
-            >>> import torch
-            >>> model = MobileCLIP("s0", device=torch.device("cpu"))
-            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
-            >>> features = model.encode_text(tokens)
         """
         try:
-            import warnings
-
-            # Suppress 'timm.models.layers is deprecated, please import via timm.layers' warning from mobileclip usage
-            with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", category=FutureWarning)
-                import mobileclip
+            import mobileclip
         except ImportError:
             # Ultralytics fork preferred since Apple MobileCLIP repo has incorrect version of torchvision
             checks.check_requirements("git+https://github.com/ultralytics/mobileclip.git")
@@ -236,8 +216,7 @@ class MobileCLIP(TextModel):
         self.eval()
 
     def tokenize(self, texts: list[str]) -> torch.Tensor:
-        """
-        Convert input texts to MobileCLIP tokens.
+        """Convert input texts to MobileCLIP tokens.
 
         Args:
             texts (list[str]): List of text strings to tokenize.
@@ -253,8 +232,7 @@ class MobileCLIP(TextModel):
 
     @smart_inference_mode()
     def encode_text(self, texts: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
-        """
-        Encode tokenized texts into normalized feature vectors.
+        """Encode tokenized texts into normalized feature vectors.
 
         Args:
             texts (torch.Tensor): Tokenized text inputs.
@@ -276,8 +254,7 @@ class MobileCLIP(TextModel):
 
 
 class MobileCLIPTS(TextModel):
-    """
-    Load a TorchScript traced version of MobileCLIP.
+    """Load a TorchScript traced version of MobileCLIP.
 
     This class implements the TextModel interface using Apple's MobileCLIP model in TorchScript format, providing
     efficient text encoding capabilities for vision-language tasks with optimized inference performance.
@@ -298,48 +275,46 @@ class MobileCLIPTS(TextModel):
         >>> features = text_encoder.encode_text(tokens)
     """
 
-    def __init__(self, device: torch.device):
-        """
-        Initialize the MobileCLIP TorchScript text encoder.
+    def __init__(self, device: torch.device, weight: str = "mobileclip_blt.ts"):
+        """Initialize the MobileCLIP TorchScript text encoder.
 
-        This class implements the TextModel interface using Apple's MobileCLIP model in TorchScript format for
-        efficient text encoding with optimized inference performance.
+        This class implements the TextModel interface using Apple's MobileCLIP model in TorchScript format for efficient
+        text encoding with optimized inference performance.
 
         Args:
             device (torch.device): Device to load the model on.
-
-        Examples:
-            >>> model = MobileCLIPTS(device=torch.device("cpu"))
-            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
-            >>> features = model.encode_text(tokens)
+            weight (str): Path to the TorchScript model weights.
         """
         super().__init__()
         from ultralytics.utils.downloads import attempt_download_asset
 
-        self.encoder = torch.jit.load(attempt_download_asset("mobileclip_blt.ts"), map_location=device)
+        self.encoder = torch.jit.load(attempt_download_asset(weight), map_location=device)
         self.tokenizer = clip.clip.tokenize
         self.device = device
 
-    def tokenize(self, texts: list[str]) -> torch.Tensor:
-        """
-        Convert input texts to MobileCLIP tokens.
+    def tokenize(self, texts: list[str], truncate: bool = True) -> torch.Tensor:
+        """Convert input texts to MobileCLIP tokens.
 
         Args:
             texts (list[str]): List of text strings to tokenize.
+            truncate (bool, optional): Whether to trim texts that exceed the tokenizer context length. Defaults to True,
+                matching CLIP's behavior to prevent runtime failures on long captions.
 
         Returns:
             (torch.Tensor): Tokenized text inputs with shape (batch_size, sequence_length).
 
         Examples:
-            >>> model = MobileCLIPTS("cpu")
+            >>> model = MobileCLIPTS(device=torch.device("cpu"))
             >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> strict_tokens = model.tokenize(
+            ...     ["a very long caption"], truncate=False
+            ... )  # RuntimeError if exceeds 77-token
         """
-        return self.tokenizer(texts).to(self.device)
+        return self.tokenizer(texts, truncate=truncate).to(self.device)
 
     @smart_inference_mode()
     def encode_text(self, texts: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
-        """
-        Encode tokenized texts into normalized feature vectors.
+        """Encode tokenized texts into normalized feature vectors.
 
         Args:
             texts (torch.Tensor): Tokenized text inputs.
@@ -360,8 +335,7 @@ class MobileCLIPTS(TextModel):
 
 
 def build_text_model(variant: str, device: torch.device = None) -> TextModel:
-    """
-    Build a text encoding model based on the specified variant.
+    """Build a text encoding model based on the specified variant.
 
     Args:
         variant (str): Model variant in format "base:size" (e.g., "clip:ViT-B/32" or "mobileclip:s0").
@@ -379,5 +353,7 @@ def build_text_model(variant: str, device: torch.device = None) -> TextModel:
         return CLIP(size, device)
     elif base == "mobileclip":
         return MobileCLIPTS(device)
+    elif base == "mobileclip2":
+        return MobileCLIPTS(device, weight="mobileclip2_b.ts")
     else:
         raise ValueError(f"Unrecognized base model: '{base}'. Supported base models: 'clip', 'mobileclip'.")

ultralytics/optim/__init__.py

@@ -0,0 +1,5 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .muon import Muon, MuSGD
+
+__all__ = ["MuSGD", "Muon"]

ultralytics/optim/muon.py

@@ -0,0 +1,338 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import torch
+from torch import optim
+
+
+def zeropower_via_newtonschulz5(G: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
+    """Compute the zeroth power / orthogonalization of matrix G using Newton-Schulz iteration.
+
+    This function implements a quintic Newton-Schulz iteration to compute an approximate orthogonalization of the input
+    matrix G. The iteration coefficients are optimized to maximize convergence slope at zero, producing a result similar
+    to UV^T from SVD, where USV^T = G, but with relaxed convergence guarantees that empirically work well for
+    optimization purposes.
+
+    Args:
+        G (torch.Tensor): Input 2D tensor/matrix to orthogonalize.
+        eps (float, optional): Small epsilon value added to norm for numerical stability. Default: 1e-7.
+
+    Returns:
+        (torch.Tensor): Orthogonalized matrix with same shape as input G.
+
+    Examples:
+        >>> G = torch.randn(128, 64)
+        >>> G_ortho = zeropower_via_newtonschulz5(G)
+        >>> print(G_ortho.shape)
+        torch.Size([128, 64])
+
+    Notes:
+        - Uses bfloat16 precision for computation.
+        - Performs exactly 5 Newton-Schulz iteration steps with fixed coefficients.
+        - Automatically transposes for efficiency when rows > columns.
+        - Output approximates US'V^T where S' has diagonal entries ~ Uniform(0.5, 1.5).
+        - Does not produce exact UV^T but works well empirically for neural network optimization.
+    """
+    assert len(G.shape) == 2
+    X = G.bfloat16()
+    X /= X.norm() + eps  # ensure top singular value <= 1
+    if G.size(0) > G.size(1):
+        X = X.T
+    for a, b, c in [  # num_steps fixed at 5
+        # original params
+        (3.4445, -4.7750, 2.0315),
+        (3.4445, -4.7750, 2.0315),
+        (3.4445, -4.7750, 2.0315),
+        (3.4445, -4.7750, 2.0315),
+        (3.4445, -4.7750, 2.0315),
+    ]:
+        # for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * A @ A
+        X = a * X + B @ X
+    if G.size(0) > G.size(1):
+        X = X.T
+    return X
+
+
+def muon_update(grad: torch.Tensor, momentum: torch.Tensor, beta: float = 0.95, nesterov: bool = True) -> torch.Tensor:
+    """Compute Muon optimizer update with momentum and orthogonalization.
+
+    This function applies momentum to the gradient, optionally uses Nesterov acceleration, and then orthogonalizes the
+    update using Newton-Schulz iterations. For convolutional filters (4D tensors), it reshapes before orthogonalization
+    and scales the final update based on parameter dimensions.
+
+    Args:
+        grad (torch.Tensor): Gradient tensor to update. Can be 2D or 4D (for conv filters).
+        momentum (torch.Tensor): Momentum buffer tensor, modified in-place via lerp.
+        beta (float, optional): Momentum coefficient for exponential moving average. Default: 0.95.
+        nesterov (bool, optional): Whether to use Nesterov momentum acceleration. Default: True.
+
+    Returns:
+        (torch.Tensor): Orthogonalized update tensor with same shape as input grad. For 4D inputs, returns reshaped
+            result matching original dimensions.
+
+    Examples:
+        >>> grad = torch.randn(64, 128)
+        >>> momentum = torch.zeros_like(grad)
+        >>> update = muon_update(grad, momentum, beta=0.95, nesterov=True)
+        >>> print(update.shape)
+        torch.Size([64, 128])
+
+    Notes:
+        - Momentum buffer is updated in-place: momentum = beta * momentum + (1-beta) * grad.
+        - With Nesterov: update = beta * momentum + (1-beta) * grad.
+        - Without Nesterov: update = momentum.
+        - 4D tensors (conv filters) are reshaped to 2D as (channels, height*width*depth) for orthogonalization.
+        - Final update is scaled by sqrt(max(dim[-2], dim[-1])) to account for parameter dimensions.
+    """
+    momentum.lerp_(grad, 1 - beta)
+    update = grad.lerp(momentum, beta) if nesterov else momentum
+    if update.ndim == 4:  # for the case of conv filters
+        update = update.view(len(update), -1)
+    update = zeropower_via_newtonschulz5(update)
+    update *= max(1, grad.size(-2) / grad.size(-1)) ** 0.5
+    return update
+
+
+class MuSGD(optim.Optimizer):
+    """Hybrid optimizer combining Muon and SGD updates for neural network training.
+
+    This optimizer implements a combination of Muon (a momentum-based optimizer with orthogonalization via Newton-Schulz
+    iterations) and standard SGD with momentum. It allows different parameter groups to use either the hybrid Muon+SGD
+    approach or pure SGD.
+
+    Args:
+        param_groups (list): List of parameter groups with their optimization settings.
+        muon (float, optional): Weight factor for Muon updates in hybrid mode. Default: 0.5.
+        sgd (float, optional): Weight factor for SGD updates in hybrid mode. Default: 0.5.
+
+    Attributes:
+        muon (float): Scaling factor applied to Muon learning rate.
+        sgd (float): Scaling factor applied to SGD learning rate in hybrid mode.
+
+    Examples:
+        >>> param_groups = [
+        ...     {
+        ...         "params": model.conv_params,
+        ...         "lr": 0.02,
+        ...         "use_muon": True,
+        ...         "momentum": 0.95,
+        ...         "nesterov": True,
+        ...         "weight_decay": 0.01,
+        ...     },
+        ...     {
+        ...         "params": model.other_params,
+        ...         "lr": 0.01,
+        ...         "use_muon": False,
+        ...         "momentum": 0.9,
+        ...         "nesterov": False,
+        ...         "weight_decay": 0,
+        ...     },
+        ... ]
+        >>> optimizer = MuSGD(param_groups, muon=0.5, sgd=0.5)
+        >>> loss = model(data)
+        >>> loss.backward()
+        >>> optimizer.step()
+
+    Notes:
+        - Parameter groups with 'use_muon': True will receive both Muon and SGD updates.
+        - Parameter groups with 'use_muon': False will receive only SGD updates.
+        - The Muon update uses orthogonalization which works best for 2D+ parameter tensors.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        momentum: float = 0.0,
+        weight_decay: float = 0.0,
+        nesterov: bool = False,
+        use_muon: bool = False,
+        muon: float = 0.5,
+        sgd: float = 0.5,
+    ):
+        """Initialize MuSGD optimizer with hybrid Muon and SGD capabilities.
+
+        Args:
+            params: Iterable of parameters to optimize or dicts defining parameter groups.
+            lr (float): Learning rate.
+            momentum (float): Momentum factor for SGD.
+            weight_decay (float): Weight decay (L2 penalty).
+            nesterov (bool): Whether to use Nesterov momentum.
+            use_muon (bool): Whether to enable Muon updates.
+            muon (float): Scaling factor for Muon component.
+            sgd (float): Scaling factor for SGD component.
+        """
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            use_muon=use_muon,
+        )
+        super().__init__(params, defaults)
+        self.muon = muon
+        self.sgd = sgd
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Perform a single optimization step.
+
+        Applies either hybrid Muon+SGD updates or pure SGD updates depending on the
+        'use_muon' flag in each parameter group. For Muon-enabled groups, parameters
+        receive both an orthogonalized Muon update and a standard SGD momentum update.
+
+        Args:
+            closure (Callable, optional): A closure that reevaluates the model
+                and returns the loss. Default: None.
+
+        Returns:
+            (torch.Tensor | None): The loss value if closure is provided, otherwise None.
+
+        Notes:
+            - Parameters with None gradients are assigned zero gradients for synchronization.
+            - Muon updates use Newton-Schulz orthogonalization and work best on 2D+ tensors.
+            - Weight decay is applied only to the SGD component in hybrid mode.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            # Muon
+            if group["use_muon"]:
+                # generate weight updates in distributed fashion
+                for p in group["params"]:
+                    lr = group["lr"]
+                    if p.grad is None:
+                        continue
+                    grad = p.grad
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+                        state["momentum_buffer_SGD"] = torch.zeros_like(p)
+
+                    update = muon_update(
+                        grad, state["momentum_buffer"], beta=group["momentum"], nesterov=group["nesterov"]
+                    )
+                    p.add_(update.reshape(p.shape), alpha=-(lr * self.muon))
+
+                    # SGD update
+                    if group["weight_decay"] != 0:
+                        grad = grad.add(p, alpha=group["weight_decay"])
+                    state["momentum_buffer_SGD"].mul_(group["momentum"]).add_(grad)
+                    sgd_update = (
+                        grad.add(state["momentum_buffer_SGD"], alpha=group["momentum"])
+                        if group["nesterov"]
+                        else state["momentum_buffer_SGD"]
+                    )
+                    p.add_(sgd_update, alpha=-(lr * self.sgd))
+            else:  # SGD
+                for p in group["params"]:
+                    lr = group["lr"]
+                    if p.grad is None:
+                        continue
+                    grad = p.grad
+                    if group["weight_decay"] != 0:
+                        grad = grad.add(p, alpha=group["weight_decay"])
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+                    state["momentum_buffer"].mul_(group["momentum"]).add_(grad)
+                    update = (
+                        grad.add(state["momentum_buffer"], alpha=group["momentum"])
+                        if group["nesterov"]
+                        else state["momentum_buffer"]
+                    )
+                    p.add_(update, alpha=-lr)
+        return loss
+
+
+class Muon(optim.Optimizer):
+    """Muon optimizer for usage in non-distributed settings.
+
+    This optimizer implements the Muon algorithm, which combines momentum-based updates with orthogonalization via
+    Newton-Schulz iterations. It applies weight decay and learning rate scaling to parameter updates.
+
+    Args:
+        params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
+        lr (float, optional): Learning rate. Default: 0.02.
+        weight_decay (float, optional): Weight decay (L2 penalty) coefficient. Default: 0.
+        momentum (float, optional): Momentum coefficient for exponential moving average. Default: 0.95.
+
+    Attributes:
+        param_groups (list): List of parameter groups with their optimization settings.
+        state (dict): Dictionary containing optimizer state for each parameter.
+
+    Examples:
+        >>> model = YourModel()
+        >>> optimizer = Muon(model.parameters(), lr=0.02, weight_decay=0.01, momentum=0.95)
+        >>> loss = model(data)
+        >>> loss.backward()
+        >>> optimizer.step()
+
+    Notes:
+        - Designed for non-distributed training environments.
+        - Uses Muon updates with orthogonalization for all parameters.
+        - Weight decay is applied multiplicatively before parameter update.
+        - Parameters with None gradients are assigned zero gradients for synchronization.
+    """
+
+    def __init__(self, params, lr: float = 0.02, weight_decay: float = 0, momentum: float = 0.95):
+        """Initialize Muon optimizer with orthogonalization-based updates.
+
+        Args:
+            params: Iterable of parameters to optimize or dicts defining parameter groups.
+            lr (float): Learning rate.
+            weight_decay (float): Weight decay factor applied multiplicatively.
+            momentum (float): Momentum factor for gradient accumulation.
+        """
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Perform a single optimization step.
+
+        Applies Muon updates to all parameters, incorporating momentum and orthogonalization.
+        Weight decay is applied multiplicatively before the parameter update.
+
+        Args:
+            closure (Callable[[], torch.Tensor] | None, optional): A closure that reevaluates the model
+                and returns the loss. Default: None.
+
+        Returns:
+            (torch.Tensor | None): The loss value if closure is provided, otherwise None.
+
+        Examples:
+            >>> optimizer = Muon(model.parameters())
+            >>> loss = model(inputs)
+            >>> loss.backward()
+            >>> optimizer.step()
+
+        Notes:
+            - Parameters with None gradients are assigned zero gradients for synchronization.
+            - Weight decay is applied as: p *= (1 - lr * weight_decay).
+            - Muon update uses Newton-Schulz orthogonalization and works best on 2D+ tensors.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    # continue
+                    p.grad = torch.zeros_like(p)  # Force synchronization
+                state = self.state[p]
+                if len(state) == 0:
+                    state["momentum_buffer"] = torch.zeros_like(p)
+                update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"])
+                p.mul_(1 - group["lr"] * group["weight_decay"])
+                p.add_(update.reshape(p.shape), alpha=-group["lr"])
+
+        return loss

ultralytics/solutions/__init__.py

@@ -19,23 +19,23 @@ from .trackzone import TrackZone
 from .vision_eye import VisionEye
 
 __all__ = (
-    "ObjectCounter",
-    "ObjectCropper",
-    "ObjectBlurrer",
     "AIGym",
-    "RegionCounter",
-    "SecurityAlarm",
+    "Analytics",
+    "DistanceCalculation",
     "Heatmap",
+    "Inference",
     "InstanceSegmentation",
-    "VisionEye",
-    "SpeedEstimator",
-    "DistanceCalculation",
-    "QueueManager",
+    "ObjectBlurrer",
+    "ObjectCounter",
+    "ObjectCropper",
     "ParkingManagement",
     "ParkingPtsSelection",
-    "Analytics",
-    "Inference",
-    "TrackZone",
+    "QueueManager",
+    "RegionCounter",
     "SearchApp",
+    "SecurityAlarm",
+    "SpeedEstimator",
+    "TrackZone",
+    "VisionEye",
     "VisualAISearch",
 )