dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/engine/tuner.py

@@ -14,57 +14,79 @@ Examples:
     >>> model.tune(data="coco8.yaml", epochs=10, iterations=300, optimizer="AdamW", plots=False, save=False, val=False)
 """
 
+from __future__ import annotations
+
+import gc
 import random
 import shutil
 import subprocess
 import time
+from datetime import datetime
 
 import numpy as np
 import torch
 
 from ultralytics.cfg import get_cfg, get_save_dir
 from ultralytics.utils import DEFAULT_CFG, LOGGER, YAML, callbacks, colorstr, remove_colorstr
+from ultralytics.utils.checks import check_requirements
+from ultralytics.utils.patches import torch_load
 from ultralytics.utils.plotting import plot_tune_results
 
 
 class Tuner:
-    """
-    A class for hyperparameter tuning of YOLO models.
+    """A class for hyperparameter tuning of YOLO models.
 
     The class evolves YOLO model hyperparameters over a given number of iterations by mutating them according to the
-    search space and retraining the model to evaluate their performance.
+    search space and retraining the model to evaluate their performance. Supports both local CSV storage and distributed
+    MongoDB Atlas coordination for multi-machine hyperparameter optimization.
 
     Attributes:
-        space (dict): Hyperparameter search space containing bounds and scaling factors for mutation.
+        space (dict[str, tuple]): Hyperparameter search space containing bounds and scaling factors for mutation.
         tune_dir (Path): Directory where evolution logs and results will be saved.
         tune_csv (Path): Path to the CSV file where evolution logs are saved.
         args (dict): Configuration arguments for the tuning process.
         callbacks (list): Callback functions to be executed during tuning.
         prefix (str): Prefix string for logging messages.
+        mongodb (MongoClient): Optional MongoDB client for distributed tuning.
+        collection (Collection): MongoDB collection for storing tuning results.
 
     Methods:
-        _mutate: Mutates the given hyperparameters within the specified bounds.
-        __call__: Executes the hyperparameter evolution across multiple iterations.
+        _mutate: Mutate hyperparameters based on bounds and scaling factors.
+        __call__: Execute the hyperparameter evolution across multiple iterations.
 
     Examples:
         Tune hyperparameters for YOLO11n on COCO8 at imgsz=640 and epochs=30 for 300 tuning iterations.
         >>> from ultralytics import YOLO
         >>> model = YOLO("yolo11n.pt")
         >>> model.tune(
-        ...     data="coco8.yaml", epochs=10, iterations=300, optimizer="AdamW", plots=False, save=False, val=False
-        ... )
+        >>>     data="coco8.yaml",
+        >>>     epochs=10,
+        >>>     iterations=300,
+        >>>     plots=False,
+        >>>     save=False,
+        >>>     val=False
+        >>> )
+
+        Tune with distributed MongoDB Atlas coordination across multiple machines:
+        >>> model.tune(
+        >>>     data="coco8.yaml",
+        >>>     epochs=10,
+        >>>     iterations=300,
+        >>>     mongodb_uri="mongodb+srv://user:pass@cluster.mongodb.net/",
+        >>>     mongodb_db="ultralytics",
+        >>>     mongodb_collection="tune_results"
+        >>> )
 
-        Tune with custom search space.
-        >>> model.tune(space={key1: val1, key2: val2})  # custom search space dictionary
+        Tune with custom search space:
+        >>> model.tune(space={"lr0": (1e-5, 1e-1), "momentum": (0.6, 0.98)})
     """
 
-    def __init__(self, args=DEFAULT_CFG, _callbacks=None):
-        """
-        Initialize the Tuner with configurations.
+    def __init__(self, args=DEFAULT_CFG, _callbacks: list | None = None):
+        """Initialize the Tuner with configurations.
 
         Args:
             args (dict): Configuration for hyperparameter evolution.
-            _callbacks (list, optional): Callback functions to be executed during tuning.
+            _callbacks (list | None, optional): Callback functions to be executed during tuning.
         """
         self.space = args.pop("space", None) or {  # key: (min, max, gain(optional))
             # 'optimizer': tune.choice(['SGD', 'Adam', 'AdamW', 'NAdam', 'RAdam', 'RMSProp']),
@@ -75,7 +97,7 @@ class Tuner:
             "warmup_epochs": (0.0, 5.0),  # warmup epochs (fractions ok)
             "warmup_momentum": (0.0, 0.95),  # warmup initial momentum
             "box": (1.0, 20.0),  # box loss gain
-            "cls": (0.2, 4.0),  # cls loss gain (scale with pixels)
+            "cls": (0.1, 4.0),  # cls loss gain (scale with pixels)
             "dfl": (0.4, 6.0),  # dfl loss gain
             "hsv_h": (0.0, 0.1),  # image HSV-Hue augmentation (fraction)
             "hsv_s": (0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
@@ -92,7 +114,12 @@ class Tuner:
             "mixup": (0.0, 1.0),  # image mixup (probability)
             "cutmix": (0.0, 1.0),  # image cutmix (probability)
             "copy_paste": (0.0, 1.0),  # segment copy-paste (probability)
+            "close_mosaic": (0.0, 10.0),  # close dataloader mosaic (epochs)
         }
+        mongodb_uri = args.pop("mongodb_uri", None)
+        mongodb_db = args.pop("mongodb_db", "ultralytics")
+        mongodb_collection = args.pop("mongodb_collection", "tuner_results")
+
         self.args = get_cfg(overrides=args)
         self.args.exist_ok = self.args.resume  # resume w/ same tune_dir
         self.tune_dir = get_save_dir(self.args, name=self.args.name or "tune")
@@ -101,88 +128,252 @@ class Tuner:
         self.callbacks = _callbacks or callbacks.get_default_callbacks()
         self.prefix = colorstr("Tuner: ")
         callbacks.add_integration_callbacks(self)
+
+        # MongoDB Atlas support (optional)
+        self.mongodb = None
+        if mongodb_uri:
+            self._init_mongodb(mongodb_uri, mongodb_db, mongodb_collection)
+
         LOGGER.info(
             f"{self.prefix}Initialized Tuner instance with 'tune_dir={self.tune_dir}'\n"
             f"{self.prefix}💡 Learn about tuning at https://docs.ultralytics.com/guides/hyperparameter-tuning"
         )
 
-    def _mutate(self, parent="single", n=5, mutation=0.8, sigma=0.2):
+    def _connect(self, uri: str = "mongodb+srv://username:password@cluster.mongodb.net/", max_retries: int = 3):
+        """Create MongoDB client with exponential backoff retry on connection failures.
+
+        Args:
+            uri (str): MongoDB connection string with credentials and cluster information.
+            max_retries (int): Maximum number of connection attempts before giving up.
+
+        Returns:
+            (MongoClient): Connected MongoDB client instance.
+        """
+        check_requirements("pymongo")
+
+        from pymongo import MongoClient
+        from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
+
+        for attempt in range(max_retries):
+            try:
+                client = MongoClient(
+                    uri,
+                    serverSelectionTimeoutMS=30000,
+                    connectTimeoutMS=20000,
+                    socketTimeoutMS=40000,
+                    retryWrites=True,
+                    retryReads=True,
+                    maxPoolSize=30,
+                    minPoolSize=3,
+                    maxIdleTimeMS=60000,
+                )
+                client.admin.command("ping")  # Test connection
+                LOGGER.info(f"{self.prefix}Connected to MongoDB Atlas (attempt {attempt + 1})")
+                return client
+            except (ConnectionFailure, ServerSelectionTimeoutError):
+                if attempt == max_retries - 1:
+                    raise
+                wait_time = 2**attempt
+                LOGGER.warning(
+                    f"{self.prefix}MongoDB connection failed (attempt {attempt + 1}), retrying in {wait_time}s..."
+                )
+                time.sleep(wait_time)
+
+    def _init_mongodb(self, mongodb_uri="", mongodb_db="", mongodb_collection=""):
+        """Initialize MongoDB connection for distributed tuning.
+
+        Connects to MongoDB Atlas for distributed hyperparameter optimization across multiple machines. Each worker
+        saves results to a shared collection and reads the latest best hyperparameters from all workers for evolution.
+
+        Args:
+            mongodb_uri (str): MongoDB connection string, e.g. 'mongodb+srv://username:password@cluster.mongodb.net/'.
+            mongodb_db (str, optional): Database name.
+            mongodb_collection (str, optional): Collection name.
+
+        Notes:
+            - Creates a fitness index for fast queries of top results
+            - Falls back to CSV-only mode if connection fails
+            - Uses connection pooling and retry logic for production reliability
+        """
+        self.mongodb = self._connect(mongodb_uri)
+        self.collection = self.mongodb[mongodb_db][mongodb_collection]
+        self.collection.create_index([("fitness", -1)], background=True)
+        LOGGER.info(f"{self.prefix}Using MongoDB Atlas for distributed tuning")
+
+    def _get_mongodb_results(self, n: int = 5) -> list:
+        """Get top N results from MongoDB sorted by fitness.
+
+        Args:
+            n (int): Number of top results to retrieve.
+
+        Returns:
+            (list[dict]): List of result documents with fitness scores and hyperparameters.
+        """
+        try:
+            return list(self.collection.find().sort("fitness", -1).limit(n))
+        except Exception:
+            return []
+
+    def _save_to_mongodb(self, fitness: float, hyperparameters: dict[str, float], metrics: dict, iteration: int):
+        """Save results to MongoDB with proper type conversion.
+
+        Args:
+            fitness (float): Fitness score achieved with these hyperparameters.
+            hyperparameters (dict[str, float]): Dictionary of hyperparameter values.
+            metrics (dict): Complete training metrics dictionary (mAP, precision, recall, losses, etc.).
+            iteration (int): Current iteration number.
+        """
+        try:
+            self.collection.insert_one(
+                {
+                    "fitness": float(fitness),
+                    "hyperparameters": {k: (v.item() if hasattr(v, "item") else v) for k, v in hyperparameters.items()},
+                    "metrics": metrics,
+                    "timestamp": datetime.now(),
+                    "iteration": iteration,
+                }
+            )
+        except Exception as e:
+            LOGGER.warning(f"{self.prefix}MongoDB save failed: {e}")
+
+    def _sync_mongodb_to_csv(self):
+        """Sync MongoDB results to CSV for plotting compatibility.
+
+        Downloads all results from MongoDB and writes them to the local CSV file in chronological order. This enables
+        the existing plotting functions to work seamlessly with distributed MongoDB data.
         """
-        Mutate hyperparameters based on bounds and scaling factors specified in `self.space`.
+        try:
+            # Get all results from MongoDB
+            all_results = list(self.collection.find().sort("iteration", 1))
+            if not all_results:
+                return
+
+            # Write to CSV
+            headers = ",".join(["fitness", *list(self.space.keys())]) + "\n"
+            with open(self.tune_csv, "w", encoding="utf-8") as f:
+                f.write(headers)
+                for result in all_results:
+                    fitness = result["fitness"]
+                    hyp_values = [result["hyperparameters"][k] for k in self.space.keys()]
+                    log_row = [round(fitness, 5), *hyp_values]
+                    f.write(",".join(map(str, log_row)) + "\n")
+
+        except Exception as e:
+            LOGGER.warning(f"{self.prefix}MongoDB to CSV sync failed: {e}")
+
+    def _crossover(self, x: np.ndarray, alpha: float = 0.2, k: int = 9) -> np.ndarray:
+        """BLX-α crossover from up to top-k parents (x[:,0]=fitness, rest=genes)."""
+        k = min(k, len(x))
+        # fitness weights (shifted to >0); fallback to uniform if degenerate
+        weights = x[:, 0] - x[:, 0].min() + 1e-6
+        if not np.isfinite(weights).all() or weights.sum() == 0:
+            weights = np.ones_like(weights)
+        idxs = random.choices(range(len(x)), weights=weights, k=k)
+        parents_mat = np.stack([x[i][1:] for i in idxs], 0)  # (k, ng) strip fitness
+        lo, hi = parents_mat.min(0), parents_mat.max(0)
+        span = hi - lo
+        return np.random.uniform(lo - alpha * span, hi + alpha * span)
+
+    def _mutate(
+        self,
+        n: int = 9,
+        mutation: float = 0.5,
+        sigma: float = 0.2,
+    ) -> dict[str, float]:
+        """Mutate hyperparameters based on bounds and scaling factors specified in `self.space`.
 
         Args:
-            parent (str): Parent selection method: 'single' or 'weighted'.
-            n (int): Number of parents to consider.
+            parent (str): Parent selection method (kept for API compatibility, unused in BLX mode).
+            n (int): Number of top parents to consider.
             mutation (float): Probability of a parameter mutation in any given iteration.
             sigma (float): Standard deviation for Gaussian random number generator.
 
         Returns:
-            (dict): A dictionary containing mutated hyperparameters.
+            (dict[str, float]): A dictionary containing mutated hyperparameters.
         """
-        if self.tune_csv.exists():  # if CSV file exists: select best hyps and mutate
-            # Select parent(s)
-            x = np.loadtxt(self.tune_csv, ndmin=2, delimiter=",", skiprows=1)
-            fitness = x[:, 0]  # first column
-            n = min(n, len(x))  # number of previous results to consider
-            x = x[np.argsort(-fitness)][:n]  # top n mutations
-            w = x[:, 0] - x[:, 0].min() + 1e-6  # weights (sum > 0)
-            if parent == "single" or len(x) == 1:
-                # x = x[random.randint(0, n - 1)]  # random selection
-                x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
-            elif parent == "weighted":
-                x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
-
-            # Mutate
-            r = np.random  # method
-            r.seed(int(time.time()))
-            g = np.array([v[2] if len(v) == 3 else 1.0 for v in self.space.values()])  # gains 0-1
+        x = None
+
+        # Try MongoDB first if available
+        if self.mongodb:
+            results = self._get_mongodb_results(n)
+            if results:
+                # MongoDB already sorted by fitness DESC, so results[0] is best
+                x = np.array([[r["fitness"]] + [r["hyperparameters"][k] for k in self.space.keys()] for r in results])
+            elif self.collection.name in self.collection.database.list_collection_names():  # Tuner started elsewhere
+                x = np.array([[0.0] + [getattr(self.args, k) for k in self.space.keys()]])
+
+        # Fall back to CSV if MongoDB unavailable or empty
+        if x is None and self.tune_csv.exists():
+            csv_data = np.loadtxt(self.tune_csv, ndmin=2, delimiter=",", skiprows=1)
+            if len(csv_data) > 0:
+                fitness = csv_data[:, 0]  # first column
+                order = np.argsort(-fitness)
+                x = csv_data[order][:n]  # top-n sorted by fitness DESC
+
+        # Mutate if we have data, otherwise use defaults
+        if x is not None:
+            np.random.seed(int(time.time()))
             ng = len(self.space)
-            v = np.ones(ng)
-            while all(v == 1):  # mutate until a change occurs (prevent duplicates)
-                v = (g * (r.random(ng) < mutation) * r.randn(ng) * r.random() * sigma + 1).clip(0.3, 3.0)
-            hyp = {k: float(x[i + 1] * v[i]) for i, k in enumerate(self.space.keys())}
+
+            # Crossover
+            genes = self._crossover(x)
+
+            # Mutation
+            gains = np.array([v[2] if len(v) == 3 else 1.0 for v in self.space.values()])  # gains 0-1
+            factors = np.ones(ng)
+            while np.all(factors == 1):  # mutate until a change occurs (prevent duplicates)
+                mask = np.random.random(ng) < mutation
+                step = np.random.randn(ng) * (sigma * gains)
+                factors = np.where(mask, np.exp(step), 1.0).clip(0.25, 4.0)
+            hyp = {k: float(genes[i] * factors[i]) for i, k in enumerate(self.space.keys())}
         else:
             hyp = {k: getattr(self.args, k) for k in self.space.keys()}
 
         # Constrain to limits
-        for k, v in self.space.items():
-            hyp[k] = max(hyp[k], v[0])  # lower limit
-            hyp[k] = min(hyp[k], v[1])  # upper limit
-            hyp[k] = round(hyp[k], 5)  # significant digits
+        for k, bounds in self.space.items():
+            hyp[k] = round(min(max(hyp[k], bounds[0]), bounds[1]), 5)
 
-        return hyp
+        # Update types
+        if "close_mosaic" in hyp:
+            hyp["close_mosaic"] = round(hyp["close_mosaic"])
 
-    def __call__(self, model=None, iterations=10, cleanup=True):
-        """
-        Execute the hyperparameter evolution process when the Tuner instance is called.
+        return hyp
 
-        This method iterates through the number of iterations, performing the following steps in each iteration:
+    def __call__(self, model=None, iterations: int = 10, cleanup: bool = True):
+        """Execute the hyperparameter evolution process when the Tuner instance is called.
 
-        1. Load the existing hyperparameters or initialize new ones.
-        2. Mutate the hyperparameters using the `mutate` method.
-        3. Train a YOLO model with the mutated hyperparameters.
-        4. Log the fitness score and mutated hyperparameters to a CSV file.
+        This method iterates through the specified number of iterations, performing the following steps:
+        1. Sync MongoDB results to CSV (if using distributed mode)
+        2. Mutate hyperparameters using the best previous results or defaults
+        3. Train a YOLO model with the mutated hyperparameters
+        4. Log fitness scores and hyperparameters to MongoDB and/or CSV
+        5. Track the best performing configuration across all iterations
 
         Args:
-            model (Model): A pre-initialized YOLO model to be used for training.
+            model (Model | None, optional): A pre-initialized YOLO model to be used for training.
             iterations (int): The number of generations to run the evolution for.
-            cleanup (bool): Whether to delete iteration weights to reduce storage space used during tuning.
-
-        Note:
-            The method utilizes the `self.tune_csv` Path object to read and log hyperparameters and fitness scores.
-            Ensure this path is set correctly in the Tuner instance.
+            cleanup (bool): Whether to delete iteration weights to reduce storage space during tuning.
         """
         t0 = time.time()
         best_save_dir, best_metrics = None, None
         (self.tune_dir / "weights").mkdir(parents=True, exist_ok=True)
+
+        # Sync MongoDB to CSV at startup for proper resume logic
+        if self.mongodb:
+            self._sync_mongodb_to_csv()
+
         start = 0
         if self.tune_csv.exists():
             x = np.loadtxt(self.tune_csv, ndmin=2, delimiter=",", skiprows=1)
             start = x.shape[0]
             LOGGER.info(f"{self.prefix}Resuming tuning run {self.tune_dir} from iteration {start + 1}...")
         for i in range(start, iterations):
+            # Linearly decay sigma from 0.2 → 0.1 over first 300 iterations
+            frac = min(i / 300.0, 1.0)
+            sigma_i = 0.2 - 0.1 * frac
+
             # Mutate hyperparameters
-            mutated_hyp = self._mutate()
+            mutated_hyp = self._mutate(sigma=sigma_i)
             LOGGER.info(f"{self.prefix}Starting iteration {i + 1}/{iterations} with hyperparameters: {mutated_hyp}")
 
             metrics = {}
@@ -195,18 +386,34 @@ class Tuner:
                 cmd = [*launch, "train", *(f"{k}={v}" for k, v in train_args.items())]
                 return_code = subprocess.run(cmd, check=True).returncode
                 ckpt_file = weights_dir / ("best.pt" if (weights_dir / "best.pt").exists() else "last.pt")
-                metrics = torch.load(ckpt_file)["train_metrics"]
+                metrics = torch_load(ckpt_file)["train_metrics"]
                 assert return_code == 0, "training failed"
 
+                # Cleanup
+                time.sleep(1)
+                gc.collect()
+                torch.cuda.empty_cache()
+
             except Exception as e:
                 LOGGER.error(f"training failure for hyperparameter tuning iteration {i + 1}\n{e}")
 
-            # Save results and mutated_hyp to CSV
+            # Save results - MongoDB takes precedence
             fitness = metrics.get("fitness", 0.0)
-            log_row = [round(fitness, 5)] + [mutated_hyp[k] for k in self.space.keys()]
-            headers = "" if self.tune_csv.exists() else (",".join(["fitness"] + list(self.space.keys())) + "\n")
-            with open(self.tune_csv, "a", encoding="utf-8") as f:
-                f.write(headers + ",".join(map(str, log_row)) + "\n")
+            if self.mongodb:
+                self._save_to_mongodb(fitness, mutated_hyp, metrics, i + 1)
+                self._sync_mongodb_to_csv()
+                total_mongo_iterations = self.collection.count_documents({})
+                if total_mongo_iterations >= iterations:
+                    LOGGER.info(
+                        f"{self.prefix}Target iterations ({iterations}) reached in MongoDB ({total_mongo_iterations}). Stopping."
+                    )
+                    break
+            else:
+                # Save to CSV only if no MongoDB
+                log_row = [round(fitness, 5)] + [mutated_hyp[k] for k in self.space.keys()]
+                headers = "" if self.tune_csv.exists() else (",".join(["fitness", *list(self.space.keys())]) + "\n")
+                with open(self.tune_csv, "a", encoding="utf-8") as f:
+                    f.write(headers + ",".join(map(str, log_row)) + "\n")
 
             # Get best results
             x = np.loadtxt(self.tune_csv, ndmin=2, delimiter=",", skiprows=1)
@@ -214,15 +421,15 @@ class Tuner:
             best_idx = fitness.argmax()
             best_is_current = best_idx == i
             if best_is_current:
-                best_save_dir = save_dir
+                best_save_dir = str(save_dir)
                 best_metrics = {k: round(v, 5) for k, v in metrics.items()}
                 for ckpt in weights_dir.glob("*.pt"):
                     shutil.copy2(ckpt, self.tune_dir / "weights")
-            elif cleanup:
-                shutil.rmtree(weights_dir, ignore_errors=True)  # remove iteration weights/ dir to reduce storage space
+            elif cleanup and best_save_dir:
+                shutil.rmtree(best_save_dir, ignore_errors=True)  # remove iteration dirs to reduce storage space
 
             # Plot tune results
-            plot_tune_results(self.tune_csv)
+            plot_tune_results(str(self.tune_csv))
 
             # Save and print tune results
             header = (
@@ -230,8 +437,7 @@ class Tuner:
                 f"{self.prefix}Results saved to {colorstr('bold', self.tune_dir)}\n"
                 f"{self.prefix}Best fitness={fitness[best_idx]} observed at iteration {best_idx + 1}\n"
                 f"{self.prefix}Best fitness metrics are {best_metrics}\n"
-                f"{self.prefix}Best fitness model is {best_save_dir}\n"
-                f"{self.prefix}Best fitness hyperparameters are printed below.\n"
+                f"{self.prefix}Best fitness model is {best_save_dir}"
             )
             LOGGER.info("\n" + header)
             data = {k: float(x[best_idx, i + 1]) for i, k in enumerate(self.space.keys())}