# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
"""
Train a model on a dataset.

Usage:
    $ yolo train model=yolo26n.pt data=coco8.yaml imgsz=640 epochs=100 batch=16
"""

from __future__ import annotations

import gc
import math
import os
import subprocess
import time
import warnings
from copy import copy, deepcopy
from datetime import datetime, timedelta
from functools import partial
from pathlib import Path

import numpy as np
import torch
from torch import distributed as dist
from torch import nn, optim

from ultralytics import __version__
from ultralytics.cfg import _YOLO_CLI_COMMAND, get_cfg, get_save_dir
from ultralytics.data.utils import check_cls_dataset, check_det_dataset, convert_ndjson_to_yolo_if_needed
from ultralytics.nn.distill_model import DistillationModel
from ultralytics.nn.tasks import load_checkpoint
from ultralytics.optim import MuSGD
from ultralytics.utils import (
    DEFAULT_CFG,
    GIT,
    LOCAL_RANK,
    LOGGER,
    RANK,
    TQDM,
    WINDOWS,
    YAML,
    callbacks,
    clean_url,
    colorstr,
    emojis,
)
from ultralytics.utils.autobatch import check_train_batch_size
from ultralytics.utils.checks import check_amp, check_file, check_imgsz, check_model_file_from_stem, print_args
from ultralytics.utils.dist import ddp_cleanup, generate_ddp_command
from ultralytics.utils.files import get_latest_run
from ultralytics.utils.patches import override_configs
from ultralytics.utils.plotting import plot_results
from ultralytics.utils.torch_utils import (
    TORCH_1_11,
    TORCH_2_0,
    TORCH_2_4,
    TORCH_2_13,
    EarlyStopping,
    ModelEMA,
    attempt_compile,
    autocast,
    convert_optimizer_state_dict_to_fp16,
    get_torch_device_backend,
    init_seeds,
    one_cycle,
    parse_device,
    prepare_qat,
    qat_state,
    restore_qat,
    select_device,
    strip_optimizer,
    strip_qat,
    torch_distributed_zero_first,
    unset_deterministic,
    unwrap_model,
)


class BaseTrainer:
    """A base class for creating trainers.

    This class provides the foundation for training YOLO models, handling the training loop, validation, checkpointing,
    and various training utilities. It supports both single-GPU and multi-GPU distributed training.

    Attributes:
        args (SimpleNamespace): Configuration for the trainer.
        validator (BaseValidator): Validator instance.
        model (nn.Module): Model instance.
        callbacks (defaultdict): Dictionary of callbacks.
        save_dir (Path): Directory to save results.
        wdir (Path): Directory to save weights.
        last (Path): Path to the last checkpoint.
        best (Path): Path to the best checkpoint.
        save_period (int): Save checkpoint every x epochs (disabled if < 1).
        batch_size (int): Batch size for training.
        epochs (int): Number of epochs to train for.
        start_epoch (int): Starting epoch for training.
        device (torch.device): Device to use for training.
        world_size (int): Number of devices used for training (0 for CPU/MPS).
        amp (bool): Whether Automatic Mixed Precision is enabled.
        scaler (torch.amp.GradScaler): Gradient scaler for AMP.
        data (dict): Dataset dictionary containing paths and metadata.
        ema (ModelEMA): EMA (Exponential Moving Average) of the model.
        resume (bool): Resume training from a checkpoint.
        lf (Callable): Learning rate scheduling function.
        optimizer (torch.optim.Optimizer): Optimizer for training.
        scheduler (torch.optim.lr_scheduler._LRScheduler): Learning rate scheduler.
        best_fitness (float): The best fitness value achieved.
        fitness (float): Current fitness value.
        loss (torch.Tensor): Current loss value.
        tloss (dict): Running mean of loss items.
        loss_names (tuple): Names of loss items, derived from the loss dict returned by the criterion on the first
            batch.
        csv (Path): Path to results CSV file.
        metrics (dict): Dictionary of metrics.
        plots (dict): Dictionary of plots.

    Methods:
        train: Execute the training process.
        validate: Run validation on the val set.
        save_model: Save model training checkpoints.
        get_dataset: Get train and validation datasets.
        setup_model: Load, create, or download model.
        build_optimizer: Construct an optimizer for the model.

    Examples:
        Initialize a task trainer (a BaseTrainer subclass) and start training
        >>> from ultralytics.models.yolo.detect import DetectionTrainer
        >>> trainer = DetectionTrainer(overrides={"model": "yolo26n.pt", "data": "coco8.yaml", "epochs": 1})
        >>> trainer.train()
    """

    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks: dict | None = None):
        """Initialize the BaseTrainer class.

        Args:
            cfg (str | dict | SimpleNamespace, optional): Path to a configuration file or configuration object.
            overrides (dict, optional): Configuration overrides.
            _callbacks (dict, optional): Dictionary of callback functions.
        """
        self.args = get_cfg(cfg, overrides)
        self.check_resume(overrides or {})
        if getattr(self.args, "augmentations", None) and not isinstance(self.args.augmentations[0], dict):
            import albumentations as A

            if any(isinstance(t, A.Lambda) for t in self.args.augmentations):  # to_dict() can't store user functions
                raise TypeError(
                    "A.Lambda augmentations can't be saved in checkpoints. Subclass A.ImageOnlyTransform or "
                    "A.DualTransform in an importable module instead."
                )
            self.args.augmentations = [A.to_dict(t) for t in self.args.augmentations]  # YAML/pickle-safe, DDP-safe
        self.args.device = parse_device(self.args.device)  # canonical string, resolves '-1' auto-selection once
        self.device = select_device(self.args.device)
        self.accelerator = get_torch_device_backend(self.device) if self.device.type not in {"cpu", "mps"} else None
        self.validator = None
        self.metrics = None
        self.plots = {}
        init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)

        # Dirs
        self.save_dir = get_save_dir(self.args)
        self.args.name = self.save_dir.name  # update name for loggers
        self.wdir = self.save_dir / "weights"  # weights dir
        if RANK in {-1, 0}:
            self.wdir.mkdir(parents=True, exist_ok=True)  # make dir
            self.args.save_dir = str(self.save_dir)
            YAML.save(self.save_dir / "args.yaml", vars(self.args))  # save run args
        self.last, self.best = self.wdir / "last.pt", self.wdir / "best.pt"  # checkpoint paths
        self.save_period = self.args.save_period

        self.batch_size = self.args.batch
        self.epochs = self.args.epochs
        self.start_epoch = 0
        if RANK == -1:
            print_args(vars(self.args))

        # Device
        if self.device.type in {"cpu", "mps"}:
            self.args.workers = 0  # faster CPU training as time dominated by inference, not dataloading

        # Callbacks - initialize early so on_pretrain_routine_start can capture original args.data
        self.callbacks = copy(_callbacks) if _callbacks else callbacks.get_default_callbacks()
        for k, v in self.callbacks.items():
            self.callbacks[k] = v.copy()

        # Device count in the launching process; distinct from utils.WORLD_SIZE set in spawned DDP workers
        if self.device.type in {"cpu", "mps"}:
            world_size = 0
        else:  # i.e. device='0', '0,1,2,3', 'npu:0', or '' auto-selecting a single GPU
            world_size = len(self.args.device.split(",")) if self.args.device else 1

        self.ddp = world_size > 1 and LOCAL_RANK == -1  # spawn DDP workers unless already one
        self.world_size = world_size
        # Run on_pretrain_routine_start before get_dataset() to capture original args.data (e.g., ul:// URIs)
        if RANK in {-1, 0} and not self.ddp:
            callbacks.add_integration_callbacks(self)
            self.run_callbacks("on_pretrain_routine_start")

        # Model and Dataset
        self.model = check_model_file_from_stem(self.args.model)  # add suffix, i.e. yolo26n -> yolo26n.pt
        with torch_distributed_zero_first(LOCAL_RANK):  # avoid auto-downloading dataset multiple times
            self.data = self.get_dataset()

        self.ema = None

        # Optimization utils init
        self.lf = None
        self.optimizer = None
        self.scheduler = None

        # Epoch level metrics
        self.best_fitness = None
        self.fitness = None
        self.loss = None
        self.tloss = None
        self.loss_names = ()
        self.csv = self.save_dir / "results.csv"
        if self.csv.exists() and not self.args.resume:
            self.csv.unlink()
        self.plot_idx = [0, 1, 2]
        self.nan_recovery_attempts = 0

    def add_callback(self, event: str, callback):
        """Append the given callback to the event's callback list."""
        self.callbacks[event].append(callback)

    def set_callback(self, event: str, callback):
        """Override the existing callbacks with the given callback for the specified event."""
        self.callbacks[event] = [callback]

    def run_callbacks(self, event: str):
        """Run all existing callbacks associated with a particular event."""
        for callback in self.callbacks.get(event, []):
            callback(self)

    def train(self):
        """Execute the training process, using DDP subprocess for multi-GPU or direct training for single-GPU."""
        # Run subprocess if DDP training, else train normally
        try:
            if self.ddp:
                # Argument checks
                if self.args.rect:
                    LOGGER.warning("'rect=True' is incompatible with Multi-GPU training, setting 'rect=False'")
                    self.args.rect = False
                if self.args.batch < 1.0:
                    raise ValueError(
                        "AutoBatch with batch<1 not supported for Multi-GPU training, "
                        f"please specify a valid batch size multiple of GPU count {self.world_size}, i.e. batch={self.world_size * 8}."
                    )

                # Command
                cmd, file = None, None
                try:
                    cmd, file = generate_ddp_command(self)
                    LOGGER.info(f"{colorstr('DDP:')} debug command {' '.join(cmd)}")
                    subprocess.run(cmd, check=True)
                finally:
                    if file is not None:
                        ddp_cleanup(self, str(file))

            else:
                self._do_train()
        finally:
            unset_deterministic()  # never leave deterministic state on, including the DDP parent and failed runs
        if not self.ddp:
            self.run_callbacks("teardown")

    def _setup_scheduler(self):
        """Initialize training learning rate scheduler."""
        if self.args.cos_lr:
            self.lf = one_cycle(1, self.args.lrf, self.epochs)  # cosine 1->hyp['lrf']
        else:
            self.lf = lambda x: max(1 - x / self.epochs, 0) * (1.0 - self.args.lrf) + self.args.lrf  # linear
        self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)

    def _get_warmup_iterations(self, num_batches):
        """Return warmup iterations, leaving at least the final epoch for regular training."""
        warmup_epochs = min(self.args.warmup_epochs, max(self.epochs - 1, 0))
        return round(warmup_epochs * num_batches) if warmup_epochs > 0 else 0

    def _setup_ddp(self):
        """Initialize and set the DistributedDataParallel parameters for training."""
        device_type = self.args.device.split(":", 1)[0]
        device_type = device_type if device_type in {"npu", "xpu"} else "cuda"
        devices = self.args.device.split(":", 1)[-1].split(",")
        index = int(devices[LOCAL_RANK])  # world_size > 1 guarantees a multi-device string
        self.device = torch.device(device_type, index)
        self.accelerator = get_torch_device_backend(self.device)
        self.accelerator.set_device(index)
        if device_type == "cuda":
            os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"  # set to enforce timeout
        elif device_type == "xpu" and not (hasattr(dist, "is_xccl_available") and dist.is_xccl_available()):
            raise RuntimeError("Multi-XPU training requires XCCL, which is not available in this PyTorch build.")
        dist.init_process_group(
            backend={"npu": "hccl", "xpu": "xccl"}.get(device_type, "nccl" if dist.is_nccl_available() else "gloo"),
            timeout=timedelta(seconds=10800),  # 3 hours
            rank=RANK,
            world_size=self.world_size,
        )

    def _build_train_pipeline(self):
        """Build dataloaders and update optimizer settings for the current batch size."""
        batch_size = self.batch_size // max(self.world_size, 1)
        self.train_loader = self.get_dataloader(
            self.data["train"], batch_size=batch_size, rank=LOCAL_RANK, mode="train"
        )
        final_batch_size = len(self.train_loader.sampler) % self.train_loader.batch_size or self.train_loader.batch_size
        min_imgsz = max(self.stride, int(self.args.imgsz * (1 - self.args.multi_scale))) // self.stride * self.stride
        if min_imgsz < 2 * self.stride and not self.train_loader.drop_last and final_batch_size == 1:
            raise ValueError(
                f"final batch=1 training at imgsz={min_imgsz} gives BatchNorm a single value per channel; "
                f"change batch, or use imgsz and multi_scale that keep every size >= {2 * self.stride}"
            )
        # Note: When training DOTA dataset, double batch size could get OOM on images with >2000 objects.
        self.test_loader = self.get_dataloader(
            self.data[self.args.split],
            batch_size=batch_size if self.args.task in {"obb", "semantic", "depth"} else batch_size * 2,
            rank=LOCAL_RANK,
            mode="val",
        )
        self.accumulate = max(round(self.args.nbs / self.batch_size), 1)  # accumulate loss before optimizing
        weight_decay = self.args.weight_decay * self.batch_size * self.accumulate / self.args.nbs  # scale weight_decay
        if self.optimizer is None:
            iterations = math.ceil(len(self.train_loader.dataset) / max(self.batch_size, self.args.nbs)) * self.epochs
            self.optimizer = self.build_optimizer(
                model=self.model,
                name=self.args.optimizer,
                lr=self.args.lr0,
                momentum=self.args.momentum,
                decay=weight_decay,
                iterations=iterations,
            )
            self._setup_scheduler()
        else:
            for group in self.optimizer.param_groups:
                if group.get("param_group") in {"weight", "muon"}:
                    group["weight_decay"] = weight_decay

    def _setup_train(self):
        """Configure model, optimizer, dataloaders, and training utilities before the training loop."""
        ckpt = self.setup_model()
        self.model = self.model.to(self.device)
        # channels_last (NHWC) is CUDA-only: lossless and Tensor-Core friendly there, but numerically wrong
        # on MPS and no benefit on CPU. Not auto-enabled on Windows, where it measured 3x slower (#26105).
        channels_last = self.args.channels_last is True or (
            self.args.channels_last is None and TORCH_1_11 and not WINDOWS
        )
        if channels_last and self.device.type == "cuda":
            self.model = self.model.to(memory_format=torch.channels_last)
        elif self.args.channels_last:
            LOGGER.warning(f"'channels_last=True' is only supported on CUDA, ignoring on '{self.device.type}'.")
        self.set_model_attributes()

        # Check imgsz
        gs = max(int(self.model.stride.max() if hasattr(self.model, "stride") else 32), 32)  # grid size (max stride)
        self.args.imgsz = check_imgsz(self.args.imgsz, stride=gs, floor=gs, max_dim=1)
        self.stride = gs  # for multiscale training

        # Prepare the final module structure before compile, DDP, and EMA wrap it.
        if self.args.quantize == 8:
            if self.resume and ckpt.get("modelopt"):
                restore_qat(self.model, ckpt["modelopt"])
            else:
                batch = max(self.batch_size // max(self.world_size, 1), 1) if self.batch_size >= 1 else 16
                with torch_distributed_zero_first(LOCAL_RANK), override_configs(self.args, {"cache": False}):
                    calibration_loader = self.get_dataloader(
                        self.data["train"], batch_size=batch, rank=-1, mode="train"
                    )
                self.model = prepare_qat(self.model, calibration_loader, self.preprocess_batch)
                del calibration_loader

        # Compile model (knowledge distillation runs the wrapped model eagerly and relies on
        # find_unused_parameters under DDP for the frozen teacher, so disable compilation when distilling)
        if self.args.distill_model is not None and self.args.compile:
            LOGGER.warning("'compile' is not supported with knowledge distillation and will be disabled.")
            self.args.compile = False
        self.model = attempt_compile(self.model, device=self.device, mode=self.args.compile)

        # Freeze layers
        freeze_list = (
            self.args.freeze
            if isinstance(self.args.freeze, list)
            else range(self.args.freeze)
            if isinstance(self.args.freeze, int)
            else []
        )
        always_freeze_names = [".dfl"]  # always freeze these layers
        freeze_layer_names = [f"model.{x}." for x in freeze_list] + always_freeze_names
        if isinstance(unwrap_model(self.model), DistillationModel):
            freeze_layer_names.append("teacher_model.")
        self.freeze_layer_names = freeze_layer_names
        for k, v in self.model.named_parameters():
            # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
            if any(x in k for x in freeze_layer_names):
                LOGGER.info(f"Freezing layer '{k}'")
                v.requires_grad = False
            elif not v.requires_grad and v.dtype.is_floating_point:  # only floating point Tensor can require gradients
                LOGGER.warning(
                    f"setting 'requires_grad=True' for frozen layer '{k}'. "
                    "See ultralytics.engine.trainer for customization of frozen layers."
                )
                v.requires_grad = True
        if not any(v.requires_grad for v in self.model.parameters()):
            raise RuntimeError(
                f"'freeze={self.args.freeze}' froze the entire model with no trainable parameters left. "
                f"Reduce 'freeze' or pass a list of specific layer indices."
            )

        # Check AMP
        self.amp = self.args.amp not in {False, "fp32"}
        self.amp = torch.tensor(self.amp).to(self.device)
        if self.amp and self.args.amp != "bf16" and RANK in {-1, 0}:  # Single-GPU and DDP
            callbacks_backup = callbacks.default_callbacks.copy()  # backup callbacks as check_amp() resets them
            self.amp = torch.tensor(check_amp(self.model), device=self.device)
            callbacks.default_callbacks = callbacks_backup  # restore callbacks
        if RANK > -1 and self.world_size > 1:  # DDP
            self.amp = self.amp.int()  # gloo errors with boolean
            dist.broadcast(self.amp, src=0)  # broadcast from rank 0 to all other ranks
        self.amp = bool(self.amp)  # as boolean
        if self.device.type == "npu":
            import torch_npu

            self.scaler = torch_npu.npu.amp.GradScaler(enabled=self.amp and self.args.amp != "bf16")
        else:
            self.scaler = (
                torch.amp.GradScaler(
                    self.device.type if self.device.type == "xpu" else "cuda",
                    enabled=self.amp and self.args.amp != "bf16",
                )
                if TORCH_2_4
                else torch.cuda.amp.GradScaler(enabled=self.amp and self.args.amp != "bf16")
            )
        # resume training would directly load DistillationModel so check here
        if self.args.distill_model is not None and not isinstance(unwrap_model(self.model), DistillationModel):
            self.model = DistillationModel(student_model=self.model, teacher_model=self.args.distill_model)
        if self.world_size > 1:
            # static_graph=True permits params used >1 time per forward (e.g. flow_model in
            # o2m+o2o pose loss branches) under torch.compile.
            ddp_kwargs = {"static_graph": bool(self.args.compile)} if TORCH_1_11 else {}
            ddp_kwargs["forward_sync_buffers" if TORCH_2_13 else "broadcast_buffers"] = False
            self.model = nn.parallel.DistributedDataParallel(
                self.model,
                device_ids=[self.device.index],
                find_unused_parameters=not bool(self.args.compile),
                **ddp_kwargs,
            )

        # Batch size
        if self.batch_size < 1 and RANK == -1:  # single-GPU only, estimate best batch size
            self.args.batch = self.batch_size = self.auto_batch()
        self._build_train_pipeline()
        self.validator = self.get_validator()
        self.set_class_weights()  # compute class weights after dataloader is ready
        self.ema = ModelEMA(self.model)  # after set_class_weights, so the copy carries them at any nesting depth
        if RANK in {-1, 0}:
            metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix="val")
            self.metrics = dict(zip(metric_keys, [0] * len(metric_keys)))
            if self.args.plots:
                self.plot_training_labels()

        self.stopper, self.stop = EarlyStopping(patience=self.args.patience), False
        self.resume_training(ckpt)
        self.scheduler.last_epoch = self.start_epoch - 1  # do not move
        self.run_callbacks("on_pretrain_routine_end")

    def _do_train(self):
        """Perform the full training loop including setup, epoch iteration, validation, and final evaluation."""
        if self.world_size > 1:
            self._setup_ddp()
        self._setup_train()

        nb = len(self.train_loader)  # number of batches
        nw = self._get_warmup_iterations(nb)
        last_opt_step = -1
        self.epoch_time = None
        self.epoch_time_start = time.time()
        self.train_time_start = time.time()
        self.run_callbacks("on_train_start")
        LOGGER.info(
            f"Using {len(self.train_loader.dataset)} train, {len(self.test_loader.dataset)} val images for "
            f"fraction={self.args.fraction} at imgsz={self.args.imgsz}\n"
            f"Using {self.train_loader.num_workers * (self.world_size or 1)} dataloader workers\n"
            f"Logging results to {colorstr('bold', self.save_dir)}\n"
            f"Starting training for " + (f"{self.args.time} hours..." if self.args.time else f"{self.epochs} epochs...")
        )
        if self.args.close_mosaic:
            base_idx = (self.epochs - self.args.close_mosaic) * nb
            self.plot_idx.extend([base_idx, base_idx + 1, base_idx + 2])
        epoch = self.start_epoch
        mosaic_closed = not self.args.close_mosaic  # close once when the run enters its final close_mosaic epochs
        self.optimizer.zero_grad()  # zero any resumed gradients to ensure stability on train start
        self._oom_retries = 0  # OOM auto-reduce counter for first epoch
        while True:
            self.epoch = epoch
            self.run_callbacks("on_train_epoch_start")
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")  # suppress 'Detected lr_scheduler.step() before optimizer.step()'
                self.scheduler.step()

            self._model_train()
            if RANK != -1:
                self.train_loader.sampler.set_epoch(epoch)
            pbar = enumerate(self.train_loader)
            # Update dataloader attributes (optional)
            if not mosaic_closed and epoch >= self.epochs - self.args.close_mosaic:
                mosaic_closed = True
                self._close_dataloader_mosaic()
                self.train_loader.reset()

            if RANK in {-1, 0}:
                if self.loss_names:
                    LOGGER.info(self.progress_string())
                pbar = TQDM(enumerate(self.train_loader), total=nb)
            self.tloss = None
            for i, batch in pbar:
                self.run_callbacks("on_train_batch_start")
                # Warmup
                ni = i + nb * epoch
                if ni < nw:
                    xi = [0, nw]  # x interp
                    self.accumulate = max(1, int(np.interp(ni, xi, [1, self.args.nbs / self.batch_size]).round()))
                    for x in self.optimizer.param_groups:
                        # Bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
                        x["lr"] = float(
                            np.interp(
                                ni,
                                xi,
                                [
                                    self.args.warmup_bias_lr if x.get("param_group") == "bias" else 0.0,
                                    x["initial_lr"] * self.lf(epoch),
                                ],
                            )
                        )
                        if "momentum" in x:
                            x["momentum"] = float(np.interp(ni, xi, [self.args.warmup_momentum, self.args.momentum]))

                # Forward
                try:
                    with autocast(torch.bfloat16 if self.args.amp == "bf16" else self.amp, device=self.device.type):
                        batch = self.preprocess_batch(batch)
                        if self.args.compile:
                            # Decouple inference and loss calculations for improved compile performance
                            preds = self.model(batch["img"])
                            loss, self.loss_items = unwrap_model(self.model).loss(batch, preds)
                        else:
                            loss, self.loss_items = self.model(batch)
                        self.loss = loss.sum()
                        if RANK != -1:
                            self.loss *= self.world_size
                        if not self.loss_names:  # derive loss names from the criterion's loss dict on first batch
                            self.loss_names = tuple(self.loss_items)
                            if RANK in {-1, 0}:
                                LOGGER.info(self.progress_string())
                                self.metrics.update(dict.fromkeys(self.label_loss_items(prefix="val"), 0.0))
                        self.tloss = (
                            self.loss_items
                            if self.tloss is None
                            else {k: (self.tloss[k] * i + v) / (i + 1) for k, v in self.loss_items.items()}
                        )

                    # Backward
                    self.scaler.scale(self.loss).backward()
                except RuntimeError as e:
                    is_oom = "out of memory" in str(e).lower()  # torch.cuda.OutOfMemoryError requires torch>=1.13
                    if not is_oom and not any(
                        s in str(e)
                        for s in (
                            "CUBLAS_STATUS_ALLOC_FAILED",
                            "CUDNN_STATUS_INTERNAL_ERROR",
                            "unable to find an engine",
                        )
                    ):
                        raise
                    if epoch > self.start_epoch or self._oom_retries >= 3 or RANK != -1:
                        raise  # only auto-reduce during first epoch on single GPU, max 3 retries
                    self._oom_retries += 1
                    old_batch = self.batch_size
                    self.args.batch = self.batch_size = max(self.batch_size // 2, 1)
                    error = f"{self.device.type.upper()} out of memory" if is_oom else "CUDA backend memory error"
                    LOGGER.warning(
                        f"{error} with batch={old_batch}. "
                        f"Reducing to batch={self.batch_size} and retrying ({self._oom_retries}/3)."
                    )
                    batch = loss = preds = None
                    self.loss = self.loss_items = self.tloss = None
                    if hasattr(self.train_loader, "close"):
                        self.train_loader.close()  # free the replaced loader's workers and prefetched batches
                    self._clear_memory()
                    self._build_train_pipeline()  # retain optimizer state across OOM retries
                    mosaic_closed = not self.args.close_mosaic  # the rebuilt loader reopened mosaic, re-arm the gate
                    self.validator.dataloader = self.test_loader  # the validator holds the pre-halving loader
                    self.scheduler.last_epoch = self.start_epoch - 1
                    nb = len(self.train_loader)
                    nw = self._get_warmup_iterations(nb)
                    last_opt_step = -1
                    self.optimizer.zero_grad()
                    break  # restart epoch loop with reduced batch size
                if ni - last_opt_step >= self.accumulate:
                    self.optimizer_step()
                    last_opt_step = ni

                    # Timed stopping
                    if self.args.time:
                        self.stop = (time.time() - self.train_time_start) > (self.args.time * 3600)
                        if RANK != -1:  # if DDP training
                            broadcast_list = [self.stop if RANK == 0 else None]
                            dist.broadcast_object_list(broadcast_list, 0)  # broadcast 'stop' to all ranks
                            self.stop = broadcast_list[0]
                        if self.stop:  # training time exceeded
                            break

                # Log
                if RANK in {-1, 0}:
                    loss_length = len(self.tloss)
                    pbar.set_description(
                        ("%11s" * 2 + "%11.4g" * (2 + loss_length))
                        % (
                            f"{epoch + 1}/{self.epochs}",
                            f"{self._get_memory():.3g}G",  # (GB) GPU memory util
                            *self.tloss.values(),  # losses
                            batch.get("cls", batch["img"]).shape[0],  # no. of instances
                            batch["img"].shape[-1],  # imgsz, i.e 640
                        )
                    )
                    if self.args.plots and ni in self.plot_idx:
                        self.plot_training_samples(batch, ni)

                self.run_callbacks("on_train_batch_end")
                if self.stop:
                    break  # allow external stop (e.g. platform cancellation) between batches
            else:
                # for/else: this block runs only when the for loop completes without break (no OOM retry)
                self._oom_retries = 0  # reset OOM counter after successful first epoch

            if self._oom_retries and not self.stop:
                continue  # OOM recovery broke the for loop, restart with reduced batch size

            if hasattr(unwrap_model(self.model).criterion, "update"):
                unwrap_model(self.model).criterion.update()

            self.lr = {f"lr/pg{ir}": x["lr"] for ir, x in enumerate(self.optimizer.param_groups)}  # for loggers

            self.run_callbacks("on_train_epoch_end")
            if RANK in {-1, 0}:
                self.ema.update_attr(self.model, include=["yaml", "nc", "args", "names", "stride", "class_weights"])

            # Validation
            final_epoch = epoch + 1 >= self.epochs
            if self.args.val or final_epoch or self.stopper.possible_stop or self.stop:
                self._clear_memory(None if self.device.type == "mps" else 0.5)  # prevent VRAM spike
                self.metrics, self.fitness = self.validate()

            # NaN recovery
            if self._handle_nan_recovery(epoch):
                last_opt_step = -1  # redo the epoch with normal step cadence, like the OOM restart
                self.optimizer.zero_grad()  # drop the corrupted pass's gradients, including NaNs still in .grad
                continue

            self.nan_recovery_attempts = 0
            if RANK in {-1, 0}:
                self.save_metrics(metrics={**self.label_loss_items(self.tloss), **self.metrics, **self.lr})
                self.stop |= self.stopper(epoch + 1, self.fitness) or final_epoch
                if self.args.time:
                    self.stop |= (time.time() - self.train_time_start) > (self.args.time * 3600)

                # Save model
                if (self.args.save or final_epoch) and self.save_model():
                    self.run_callbacks("on_model_save")

            # Scheduler
            t = time.time()
            self.epoch_time = t - self.epoch_time_start
            self.epoch_time_start = t
            if self.args.time:
                mean_epoch_time = (t - self.train_time_start) / (epoch - self.start_epoch + 1)
                self.epochs = self.args.epochs = math.ceil(self.args.time * 3600 / mean_epoch_time)
                nw = self._get_warmup_iterations(nb)
                self._setup_scheduler()
                self.scheduler.last_epoch = self.epoch  # do not move
                self.stop |= epoch >= self.epochs  # stop if exceeded epochs
            self.run_callbacks("on_fit_epoch_end")
            # clear if memory utilization > 50%; always clear on MPS due to leak https://github.com/ultralytics/ultralytics/issues/22621
            self._clear_memory(None if self.device.type == "mps" else 0.5)

            # Early Stopping
            if RANK != -1:  # if DDP training
                broadcast_list = [self.stop if RANK == 0 else None]
                dist.broadcast_object_list(broadcast_list, 0)  # broadcast 'stop' to all ranks
                self.stop = broadcast_list[0]
            if self.stop:
                break  # must break all DDP ranks
            epoch += 1

        seconds = time.time() - self.train_time_start
        LOGGER.info(f"\n{epoch - self.start_epoch + 1} epochs completed in {seconds / 3600:.3f} hours.")
        # Do final val with best.pt
        self.final_eval()
        if RANK in {-1, 0}:
            if self.args.plots:
                self.plot_metrics()
            self.run_callbacks("on_train_end")
        self._clear_memory()
        for loader in (self.train_loader, self.test_loader):
            if hasattr(loader, "close"):
                loader.close()  # shut down persistent dataloader workers so none survive to interpreter exit

    def auto_batch(self, max_num_obj=0, dataset_size=0):
        """Calculate optimal batch size based on model and device memory constraints."""
        # Stride-aligned to match the true multi-scale max size; pyramid heads require stride-multiple inputs
        max_imgsz = math.ceil(self.args.imgsz * (1 + self.args.multi_scale) / self.stride) * self.stride
        return check_train_batch_size(
            model=self.model,
            imgsz=max_imgsz,
            amp=torch.bfloat16 if self.args.amp == "bf16" else self.amp,
            batch=self.batch_size,
            max_num_obj=max_num_obj,
            dataset_size=dataset_size,
        )  # returns batch size

    def _get_memory(self, fraction=False):
        """Get accelerator memory utilization in GB or as a fraction of total memory."""
        memory, total = 0, 0
        if self.device.type == "mps":
            memory = torch.mps.driver_allocated_memory()
            if fraction:
                return __import__("psutil").virtual_memory().percent / 100
        elif self.device.type != "cpu":
            memory = self.accelerator.memory_reserved()
            if fraction:
                total = self.accelerator.get_device_properties(self.device).total_memory
        return ((memory / total) if total > 0 else 0) if fraction else (memory / 2**30)

    def _clear_memory(self, threshold: float | None = None):
        """Clear accelerator memory by calling garbage collector and emptying cache."""
        if threshold:
            assert 0 <= threshold <= 1, "Threshold must be between 0 and 1."
            if self._get_memory(fraction=True) <= threshold:
                return
        gc.collect()
        if self.device.type == "mps":
            torch.mps.empty_cache()
        elif self.device.type == "cpu":
            return
        else:
            self.accelerator.empty_cache()

    def read_results_csv(self):
        """Read results.csv into a dictionary using polars."""
        import polars as pl  # scope for faster 'import ultralytics'

        try:
            return pl.read_csv(self.csv.read_bytes(), infer_schema_length=None).to_dict(as_series=False)
        except Exception:
            return {}

    def _model_train(self):
        """Set model in training mode."""
        self.model.train()
        # Freeze BN stat
        for n, m in self.model.named_modules():
            if any(filter(lambda f: f in n, self.freeze_layer_names)) and isinstance(m, nn.BatchNorm2d):
                m.eval()

    def save_model(self):
        """Save model training checkpoints with additional metadata.

        Returns:
            (bool): True once the checkpoints have been written.
        """
        import io

        # A transient NaN/Inf permanently poisons the EMA running average (ema = decay*ema + (1-decay)*model), so
        # save_model would otherwise skip every epoch and the run would finish with no checkpoint on valid input.
        # Resync each poisoned EMA tensor from the live model where finite; any tensor that is non-finite in both is
        # left for the nan_to_num_ pass below, so a usable checkpoint is always written.
        ema = self.ema.ema
        if not all(torch.isfinite(v).all() for v in ema.state_dict().values() if isinstance(v, torch.Tensor)):
            model_sd = unwrap_model(self.model).state_dict()
            for k, v in ema.state_dict().items():
                if isinstance(v, torch.Tensor) and not torch.isfinite(v).all() and torch.isfinite(model_sd[k]).all():
                    v.copy_(model_sd[k])
        # Serialize NCHW regardless of channels_last training: released versions fuse with .view(), which crashes on
        # NHWC-strided checkpoint weights, and trainer/predictor re-apply channels_last at setup anyway.
        modelopt = qat_state(ema)
        ema = deepcopy(ema).half().to(memory_format=torch.contiguous_format)
        if hasattr(ema, "criterion"):
            ema.criterion = None  # strip training-only state from the serialization snapshot
        # Clamp fp16 serialization overflow without mutating the live EMA.
        for v in ema.state_dict().values():
            if isinstance(v, torch.Tensor) and v.is_floating_point():
                torch.nan_to_num_(v)

        # Serialize ckpt to a byte buffer once (faster than repeated torch.save() calls)
        strip_qat(ema)
        buffer = io.BytesIO()
        torch.save(
            {
                "epoch": self.epoch,
                "best_fitness": self.best_fitness,
                "stopper": {"best_fitness": self.stopper.best_fitness, "best_epoch": self.stopper.best_epoch},
                "model": None,  # resume and final checkpoints derive from EMA
                "ema": ema,
                "updates": self.ema.updates,
                "modelopt": modelopt,  # quantization state of a QAT model, restored by load_checkpoint()
                "optimizer": convert_optimizer_state_dict_to_fp16(deepcopy(self.optimizer.state_dict())),
                "scaler": self.scaler.state_dict(),
                "train_args": vars(self.args),  # save as dict
                "train_metrics": {**self.metrics, "fitness": self.fitness},
                "train_results": self.read_results_csv(),
                "date": datetime.now().astimezone().isoformat(),
                "version": __version__,
                "git": {
                    "root": str(GIT.root),
                    "branch": GIT.branch,
                    "commit": GIT.commit,
                    "message": GIT.message,
                    "origin": GIT.origin,
                },
                "license": "AGPL-3.0 (https://ultralytics.com/license)",
                "docs": "https://docs.ultralytics.com",
            },
            buffer,
        )
        serialized_ckpt = buffer.getvalue()  # get the serialized content to save

        # Save checkpoints
        self.wdir.mkdir(parents=True, exist_ok=True)  # ensure weights directory exists
        self.last.write_bytes(serialized_ckpt)  # save last.pt
        if self.best_fitness == self.fitness:
            self.best.write_bytes(serialized_ckpt)  # save best.pt
        if (self.save_period > 0) and (self.epoch % self.save_period == 0):
            (self.wdir / f"epoch{self.epoch}.pt").write_bytes(serialized_ckpt)  # save epoch, i.e. 'epoch3.pt'
        return True

    def get_dataset(self):
        """Get train and validation datasets from data dictionary.

        Returns:
            (dict): A dictionary containing the training/validation/test dataset and category names.

        Raises:
            RuntimeError: If the dataset cannot be found or checked.
        """
        try:
            self.args.data = convert_ndjson_to_yolo_if_needed(self.args.data, self.args.fraction, split=self.args.split)

            # Task-specific dataset checking
            if self.args.task == "classify":
                data = check_cls_dataset(self.args.data, split=self.args.split)
            elif str(self.args.data).rsplit(".", 1)[-1] in {"yaml", "yml"} or self.args.task in {
                "detect",
                "segment",
                "pose",
                "obb",
                "semantic",
                "depth",
            }:
                data = check_det_dataset(self.args.data, split=self.args.split)
                if "yaml_file" in data:
                    self.args.data = data["yaml_file"]  # for validating 'yolo train data=url.zip' usage
        except Exception as e:
            raise RuntimeError(emojis(f"Dataset '{clean_url(self.args.data)}' error ❌ {e}")) from e
        if self.args.single_cls:
            LOGGER.info("Overriding class names with single class.")
            data["names"] = {0: "item"}
            data["nc"] = 1
        return data

    def setup_model(self):
        """Load, create, or download model for any task.

        Returns:
            (dict | None): Checkpoint to resume training from, or None if no checkpoint is loaded.
        """
        if isinstance(self.model, torch.nn.Module):  # if model is loaded beforehand. No setup needed
            return

        cfg, weights = self.model, None
        ckpt = None
        if str(self.model).endswith(".pt"):
            weights, ckpt = load_checkpoint(self.model)
            cfg = weights.yaml
        if isinstance(self.args.pretrained, (str, Path)) and not self.resume:
            weights, _ = load_checkpoint(self.args.pretrained)
        elif self.args.pretrained is False and not self.resume:
            weights = None

        # rebuild DistillationModel from resuming checkpoint
        if isinstance(weights, DistillationModel):
            if RANK in {-1, 0}:
                LOGGER.info("Resuming training DistillationModel from checkpoint weights")
            student_model = self.get_model(cfg=cfg, weights=weights.student_model, verbose=RANK in {-1, 0})
            student_model.args = self.args
            # teacher is stripped from the checkpoint to save memory/disk; rebuild it from the distill_model path
            teacher_model = weights.teacher_model if weights.teacher_model is not None else self.args.distill_model
            model = DistillationModel(student_model=student_model, teacher_model=teacher_model)
            if getattr(weights, "projector", None) is not None:
                model.projector.load_state_dict(weights.projector.state_dict())  # restore the trained projector
            model.criterion = None
            self.model = model
        else:
            self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK in {-1, 0})  # calls Model(cfg, weights)
        return ckpt

    def optimizer_step(self):
        """Perform a single step of the training optimizer with gradient clipping and EMA update."""
        self.scaler.unscale_(self.optimizer)  # unscale gradients
        if self.device.type == "npu" and TORCH_2_0:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0, foreach=False)
        else:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
        self.scaler.step(self.optimizer)
        self.scaler.update()
        self.optimizer.zero_grad()
        if self.ema:
            self.ema.update(self.model)

    def preprocess_batch(self, batch):
        """Allow custom preprocessing of model inputs and ground truths depending on task type."""
        return batch

    def validate(self):
        """Run validation on val set using self.validator.

        Returns:
            (tuple): A tuple containing:
                - metrics (dict | None): Dictionary of validation metrics, or None if validation was skipped.
                - fitness (float | None): Fitness score for the validation, or None if validation was skipped.
        """
        if self.ema and self.world_size > 1:
            # Sync EMA buffers from rank 0 to all ranks
            for buffer in self.ema.ema.buffers():
                dist.broadcast(buffer, src=0)
        metrics = self.validator(self)
        if metrics is None:
            return None, None
        fitness = metrics.pop("fitness", -self.loss.detach().cpu().numpy())  # use loss as fitness measure if not found
        if self.best_fitness is None or self.best_fitness < fitness:
            self.best_fitness = fitness
        return metrics, fitness

    def get_model(self, cfg=None, weights=None, verbose=True):
        """Raise NotImplementedError (must return a model built from cfg and weights in subclasses)."""
        raise NotImplementedError("This task trainer doesn't support loading cfg files")

    def get_validator(self):
        """Raise NotImplementedError (must be implemented by subclasses)."""
        raise NotImplementedError("get_validator function not implemented in trainer")

    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
        """Raise NotImplementedError (must return a `torch.utils.data.DataLoader` in subclasses)."""
        raise NotImplementedError("get_dataloader function not implemented in trainer")

    def build_dataset(self, img_path, mode="train", batch=None):
        """Build dataset."""
        raise NotImplementedError("build_dataset function not implemented in trainer")

    def label_loss_items(self, loss_items=None, prefix="train"):
        """Return a loss dict with labeled training loss items, or a list of loss names if loss_items is None."""
        if loss_items is None:
            return [f"{prefix}/{x}" for x in self.loss_names]
        return {f"{prefix}/{k}": round(float(v), 5) for k, v in loss_items.items()}

    def set_model_attributes(self):
        """Set or update model parameters before training."""
        self.model.names = self.data["names"]

    def set_class_weights(self):
        """Compute and set class weights for handling class imbalance. Override in subclasses."""

    def progress_string(self):
        """Return a string describing training progress."""
        return ""

    # TODO: may need to put these following functions into callback
    def plot_training_samples(self, batch, ni):
        """Plot training samples during YOLO training."""

    def plot_training_labels(self):
        """Plot training labels for YOLO model."""

    def save_metrics(self, metrics):
        """Save training metrics to a CSV file."""
        keys, vals = list(metrics.keys()), list(metrics.values())
        n = len(metrics) + 2  # number of cols
        t = time.time() - self.train_time_start
        self.csv.parent.mkdir(parents=True, exist_ok=True)  # ensure parent directory exists
        s = "" if self.csv.exists() else ("%s," * n % ("epoch", "time", *keys)).rstrip(",") + "\n"
        with open(self.csv, "a", encoding="utf-8") as f:
            f.write(s + ("%.6g," * n % (self.epoch + 1, t, *vals)).rstrip(",") + "\n")

    def plot_metrics(self):
        """Plot metrics from a CSV file."""
        plot_results(file=self.csv, on_plot=self.on_plot)  # save results.png

    def on_plot(self, name, data=None):
        """Register plots (e.g. to be consumed in callbacks)."""
        path = Path(name)
        self.plots[path] = {"data": data, "timestamp": time.time()}

    def final_eval(self):
        """Perform final evaluation and validation for the YOLO model."""
        model = self.best if self.best.exists() else None
        with torch_distributed_zero_first(LOCAL_RANK):  # strip only on GPU 0; other GPUs should wait
            if RANK in {-1, 0}:
                ckpt = strip_optimizer(self.last) if self.last.exists() else {}
                if model:
                    # update best.pt train_metrics from last.pt
                    strip_optimizer(self.best, updates={"train_results": ckpt.get("train_results")})
        if model:
            LOGGER.info(f"\nValidating {model}...")
            self.validator.args.plots = self.args.plots
            self.validator.args.compile = False  # disable final val compile as too slow
            self.metrics = self.validator(model=model)
            self.metrics.pop("fitness", None)
            self.epoch += 1  # log best metrics at step epochs+1, not overwriting last epoch
            self.run_callbacks("on_fit_epoch_end")
            self.epoch -= 1  # restore epoch

    def check_resume(self, overrides):
        """Check if resume checkpoint exists and update arguments accordingly."""
        resume = self.args.resume
        if resume:
            try:
                exists = isinstance(resume, (str, Path)) and Path(resume).exists()
                last = Path(check_file(resume) if exists else get_latest_run())
                ckpt_args = load_checkpoint(last)[0].args
            except Exception as e:
                raise FileNotFoundError(
                    "Resume checkpoint not found. Please pass a valid checkpoint to resume from, "
                    "i.e. 'yolo train resume model=path/to/last.pt'"
                ) from e
            if self.args.data:
                ckpt_args["data"] = self.args.data

            resume = True
            self.args = get_cfg(ckpt_args)
            self.args.model = self.args.resume = str(last)  # reinstate model
            allowed = {  # allow arg updates to reduce memory or update device on resume
                "imgsz",
                "batch",
                "device",
                "close_mosaic",
                "augmentations",
                "save_period",
                "workers",
                "cache",
                "patience",
                "time",
                "freeze",
                "val",
                "plots",
                "channels_last",
                "distill_model",
                "save_dir",
            }
            ignored = []
            for k, v in overrides.items():
                if k in allowed:
                    setattr(self.args, k, v)
                elif k not in {"model", "data", "mode", "resume", "pretrained"} and v != getattr(self.args, k, None):
                    ignored.append(k)
            if ignored:
                LOGGER.warning(f"Resume ignores {ignored}, using checkpoint values. Start a new run to change them.")
        self.resume = resume

    def _load_checkpoint_state(self, ckpt):
        """Load optimizer, scaler, EMA, best_fitness, and early stopping state from checkpoint."""
        if ckpt.get("optimizer") is not None:
            for saved, group in zip(ckpt["optimizer"]["param_groups"], self.optimizer.param_groups):
                saved["fused"] = group.get("fused")  # runtime device, not the checkpoint, picks the kernel
            self.optimizer.load_state_dict(ckpt["optimizer"])
        if ckpt.get("scaler"):
            self.scaler.load_state_dict(ckpt["scaler"])
        if self.ema and ckpt.get("ema"):
            self.ema = ModelEMA(self.model)  # validation with EMA creates inference tensors that can't be updated
            # A QAT checkpoint serializes its EMA without quantizers, but load_checkpoint() re-applied them in place
            # to this very module, so the strict load below still matches the QAT-structured EMA built above.
            self.ema.ema.load_state_dict(ckpt["ema"].float().state_dict())
            self.ema.updates = ckpt["updates"]
        self.best_fitness = ckpt.get("best_fitness")
        self.stopper.__dict__.update(ckpt.get("stopper") or {})  # older checkpoints keep a fresh stopper

    def _handle_nan_recovery(self, epoch):
        """Detect and recover from NaN/Inf loss by loading last checkpoint."""
        loss_nan = self.loss is not None and not self.loss.isfinite()
        fitness_nan = self.fitness is not None and not np.isfinite(self.fitness)
        corrupted = RANK in {-1, 0} and (loss_nan or fitness_nan)
        reason = "Loss NaN/Inf" if loss_nan else "Fitness NaN/Inf"
        if RANK != -1:  # DDP: broadcast to all ranks
            broadcast_list = [corrupted if RANK == 0 else None]
            dist.broadcast_object_list(broadcast_list, 0)
            corrupted = broadcast_list[0]
        if not corrupted:
            return False
        if epoch == self.start_epoch:
            LOGGER.warning(f"{reason} detected but can not recover from last.pt...")
            return False  # Cannot recover on first epoch, let training continue
        if not self.last.exists():
            raise RuntimeError(f"{reason} detected but no valid last.pt is available for recovery")
        self.nan_recovery_attempts += 1
        if self.nan_recovery_attempts > 3:
            raise RuntimeError(f"Training failed: NaN persisted for {self.nan_recovery_attempts} epochs")
        LOGGER.warning(f"{reason} detected (attempt {self.nan_recovery_attempts}/3), recovering from last.pt...")
        self._model_train()  # set model to train mode before loading checkpoint to avoid inference tensor errors
        _, ckpt = load_checkpoint(self.last)
        ema = ckpt["ema"].float()
        ema_state = ema.state_dict()
        if not all(torch.isfinite(v).all() for v in ema_state.values() if isinstance(v, torch.Tensor)):
            raise RuntimeError(f"Checkpoint {self.last} is corrupted with NaN/Inf weights")
        model = unwrap_model(self.model)
        if hasattr(model, "student_model"):
            # Distillation: the EMA is stripped of the teacher (rebuilt from the distill_model path), so only the
            # student and projector are restored; loading them separately keeps a strict key match.
            model.student_model.load_state_dict(ema.student_model.state_dict())
            model.projector.load_state_dict(ema.projector.state_dict())
        else:
            model.load_state_dict(ema_state)  # Load EMA weights into model
        self._load_checkpoint_state(ckpt)  # Load optimizer/scaler/EMA/best_fitness
        del ckpt, ema, ema_state
        self.scheduler.last_epoch = epoch - 1
        return True

    def resume_training(self, ckpt):
        """Resume YOLO training from a given checkpoint."""
        if ckpt is None or not self.resume:
            return
        start_epoch = ckpt.get("epoch", -1) + 1
        assert 0 < start_epoch < self.epochs, (
            f"{self.args.model} training to {self.epochs} epochs is finished, nothing to resume.\n"
            f"Start a new training without resuming, i.e. 'yolo train model={self.args.model}'"
        )
        LOGGER.info(f"Resuming training {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs")
        self._load_checkpoint_state(ckpt)
        model = unwrap_model(self.model)
        if getattr(getattr(model, "student_model", model).model[-1], "one2one_cv2", None) is not None:
            # Resume both head losses independently of the selected inference head.
            model.criterion = model.init_criterion()
            model.criterion.updates = start_epoch - 1
            model.criterion.update()
        self.start_epoch = start_epoch

    def _close_dataloader_mosaic(self):
        """Update dataloaders to stop using mosaic augmentation."""
        if hasattr(self.train_loader.dataset, "mosaic"):
            self.train_loader.dataset.mosaic = False
        if hasattr(self.train_loader.dataset, "close_mosaic"):
            LOGGER.info("Closing dataloader mosaic")
            self.train_loader.dataset.close_mosaic(hyp=copy(self.args))

    def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
        """Construct an optimizer for the given model.

        Args:
            model (torch.nn.Module): The model for which to build an optimizer.
            name (str, optional): The name of the optimizer to use. If 'auto', the optimizer is selected based on the
                number of iterations.
            lr (float, optional): The learning rate for the optimizer.
            momentum (float, optional): The momentum factor for the optimizer.
            decay (float, optional): The weight decay for the optimizer.
            iterations (float, optional): The number of iterations, which determines the optimizer if name is 'auto'.

        Returns:
            (torch.optim.Optimizer): The constructed optimizer.

        Raises:
            NotImplementedError: If the optimizer name is not supported.
        """
        g = [{}, {}, {}, {}]  # optimizer parameter groups
        bn = tuple(v for k, v in nn.__dict__.items() if "Norm" in k)  # normalization layers, i.e. BatchNorm2d()
        optimizers = {"Adam", "Adamax", "AdamW", "NAdam", "RAdam", "RMSprop", "SGD", "MuSGD", "auto"}
        name = {x.lower(): x for x in optimizers}.get(str(name).lower(), str(name))
        if name == "auto":
            LOGGER.info(
                f"{colorstr('optimizer:')} 'optimizer=auto' found, "
                f"ignoring 'lr0={self.args.lr0}' and determining best 'optimizer' and 'lr0' automatically... "
            )
            nc = self.data.get("nc", 10)  # number of classes
            lr_fit = round(0.002 * 5 / (4 + nc), 6)  # lr0 fit equation to 6 decimal places
            name, lr, momentum = ("MuSGD", 0.01, 0.9) if iterations > 10000 else ("AdamW", lr_fit, 0.9)
            self.args.optimizer, self.args.lr0 = name, lr  # resume rebuilds this choice from train_args
            self.args.warmup_bias_lr = 0.0  # no higher than 0.01 for Adam

        use_muon = name == "MuSGD"
        for module_name, module in unwrap_model(model).named_modules():
            for param_name, param in module.named_parameters(recurse=False):
                fullname = f"{module_name}.{param_name}" if module_name else param_name
                if param.ndim in {2, 4} and use_muon:  # muon only orthogonalizes matrices and conv filters
                    g[3][fullname] = param  # muon params
                elif "bias" in fullname:  # bias (no decay)
                    g[2][fullname] = param
                elif isinstance(module, bn) or "logit_scale" in fullname:  # weight (no decay)
                    # ContrastiveHead and BNContrastiveHead included here with 'logit_scale'
                    g[1][fullname] = param
                else:  # weight (with decay)
                    g[0][fullname] = param
        if not use_muon:
            g = [x.values() for x in g[:3]]  # convert to list of params

        if name in {"Adam", "Adamax", "AdamW", "NAdam", "RAdam"}:
            optim_args = {"lr": lr, "betas": (momentum, 0.999), "weight_decay": 0.0}
        elif name == "RMSprop":
            optim_args = {"lr": lr, "momentum": momentum}
        elif name == "SGD" or name == "MuSGD":
            optim_args = {"lr": lr, "momentum": momentum, "nesterov": True}
        else:
            raise NotImplementedError(
                f"Optimizer '{name}' not found in list of available optimizers {optimizers}. "
                "Request support for additional optimizers at https://github.com/ultralytics/ultralytics."
            )

        num_params = [len(g[0]), len(g[1]), len(g[2])]  # number of param groups
        g[2] = {"params": g[2], **optim_args, "param_group": "bias"}
        g[0] = {"params": g[0], **optim_args, "weight_decay": decay, "param_group": "weight"}
        g[1] = {"params": g[1], **optim_args, "weight_decay": 0.0, "param_group": "bn"}
        muon, sgd = (0.2, 1.0)
        if use_muon:
            num_params[0] = len(g[3])  # update number of params
            g[3] = {"params": g[3], **optim_args, "weight_decay": decay, "use_muon": True, "param_group": "muon"}
            # higher lr for certain parameters in MuSGD when finetuning
            target = unwrap_model(model)
            head = getattr(target, "student_model", target).model[-1]
            heads = (getattr(head, "cv3", None), getattr(head, "one2one_cv3", None))
            boosted = {id(p) for m in heads if m for p in m.parameters()}
            g_ = []  # new param groups
            for x in g:
                p = x.pop("params")
                p1, p2 = [], []
                for k, v in p.items():
                    (p1 if id(v) in boosted or "proto.semseg" in k or "SemanticSegment" in k else p2).append(v)
                g_.extend([{"params": p1, **x, "lr": lr * 3}, {"params": p2, **x}])
            g = g_
        # fused=True must go to the constructor: that is where Adam registers _step_supports_amp_scaling, which lets
        # GradScaler pass found_inf to the kernel instead of reading it back on the host
        fused = {"fused": True} if name in {"Adam", "AdamW"} and TORCH_2_0 and next(model.parameters()).is_cuda else {}
        optimizer = (partial(MuSGD, muon=muon, sgd=sgd) if use_muon else getattr(optim, name))(params=g, **fused)

        LOGGER.info(
            f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}, momentum={momentum}) with parameter groups "
            f"{num_params[1]} weight(decay=0.0), {num_params[0]} weight(decay={decay}), {num_params[2]} bias(decay=0.0)"
        )
        return optimizer


class MultiTrainer:
    """Fine-tune a single base model across a collection of datasets and aggregate per-dataset results.

    Used automatically by Model.train() when `data` is a list or tuple, allowing one base model to be benchmarked across
    many datasets (such as the RF100 collection) in a single call. The datasets are fine-tuned in series and the same
    base weights seed each run, so every run starts from an identical model. All output is grouped under one sweep
    directory (e.g. runs/detect/multitrain): each dataset gets its own run subdirectory, and the per-dataset and mean
    metrics are written to multitrain_results.json (for post-processing) alongside a multitrain_results.png bar
    chart. The base model object is left unchanged; each dataset's fine-tuned weights live in its own run directory.

    Attributes:
        trainer (type[BaseTrainer] | None): Task trainer class for Python runs, or None for CLI subprocess runs.
        args (dict): Training arguments shared across datasets; its `data` key holds the dataset collection.
        model (torch.nn.Module): Base model whose weights seed each per-dataset fine-tune.
        callbacks (dict | None): Callbacks forwarded to each per-dataset trainer.
        trainers (list[SimpleNamespace]): Completed per-dataset run records.
        metrics (dict): Mapping of each run name (e.g. coco8, coco8-2) to its training-metrics dict from the checkpoint.
        mean_metrics (dict): Mean training metrics across successful datasets.
        save_dir (Path | None): Sweep directory holding the per-dataset runs and the results JSON/plot.

    Examples:
        Fine-tune one base model across several datasets and read back per-run metrics:
        >>> from ultralytics import YOLO
        >>> model = YOLO("yolo26n.pt")
        >>> results = model.train(data=["coco8.yaml", "african-wildlife.yaml"], epochs=10)
        >>> results["coco8"]["fitness"]  # final fitness on the coco8 run
    """

    def __init__(self, trainer, args, model, _callbacks: dict | None = None):
        """Initialize MultiTrainer with a task trainer class, shared training arguments, and the base model.

        Args:
            trainer (type[BaseTrainer] | None): Task trainer class to run once per dataset. None uses CLI subprocesses.
            args (dict): Training arguments; the `data` key holds the list/tuple of datasets to fine-tune on.
            model (torch.nn.Module): Base model whose weights seed each per-dataset fine-tune.
            _callbacks (dict, optional): Callback functions forwarded to each per-dataset trainer.
        """
        self.trainer = trainer
        self.args = args
        self.model = model
        self.callbacks = _callbacks
        self.trainers = []
        self.metrics = {}
        self.mean_metrics = {}
        self.save_dir = None

    def train(self):
        """Fine-tune the base model on each dataset in series and return a {dataset: metrics} mapping."""
        from types import SimpleNamespace

        from ultralytics.utils.patches import torch_load, torch_save

        datasets = self.args["data"]
        # Group every per-dataset run and the summary plot under one sweep directory, e.g. runs/detect/multitrain
        sweep = SimpleNamespace(
            project=self.args.get("project"),
            task=self.args.get("task"),
            mode="train",
            exist_ok=self.args.get("exist_ok", False),
        )
        self.save_dir = get_save_dir(sweep, name="multitrain")
        self.save_dir.mkdir(parents=True, exist_ok=True)
        model_name = Path(str(self.args.get("model") or "multitrain_base")).stem
        base_model = self.save_dir / f"{model_name}.pt" if self.trainer is None else None
        if base_model:
            state = qat_state(self.model)
            model = deepcopy(self.model).half()
            strip_qat(model)
            torch_save({"model": model, "modelopt": state, "train_args": getattr(self.model, "args", {})}, base_model)
        try:
            for i, data in enumerate(datasets):
                LOGGER.info(
                    f"\n{colorstr('blue', 'bold', f'MultiTrainer {i + 1}/{len(datasets)}:')} fine-tuning on {data}"
                )
                path = Path(str(data))
                parent = path.parent.name
                name = Path(os.path.abspath(path.parent)).name if path.stem == "data" and parent else path.stem
                run_name = name
                try:
                    overrides = {
                        **self.args,
                        "data": data,
                        "project": str(self.save_dir),  # nest per-dataset runs inside the sweep directory
                        "name": name,
                        "resume": False,
                    }
                    run = SimpleNamespace(
                        project=overrides["project"],
                        name=overrides["name"],
                        task=overrides.get("task"),
                        mode="train",
                        exist_ok=overrides.get("exist_ok", False),
                        save_dir=None,
                    )
                    save_dir = get_save_dir(run)
                    save_dir.mkdir(parents=True, exist_ok=True)
                    run_name = save_dir.name
                    overrides["save_dir"] = str(save_dir)
                    if self.trainer is None:
                        overrides["model"] = str(base_model)
                        overrides.pop("cfg", None)  # already merged here; the CLI would re-apply the file
                        subprocess.run(
                            [
                                *_YOLO_CLI_COMMAND,
                                "train",
                                *(f"{k}={v}" for k, v in overrides.items()),
                            ],
                            check=True,
                        )
                    else:
                        trainer = self.trainer(overrides=overrides, _callbacks=self.callbacks)
                        pretrained = overrides.get("pretrained", True)
                        weights = None if pretrained is False else self.model
                        if isinstance(pretrained, (str, Path)):
                            weights, _ = load_checkpoint(pretrained)
                        trainer.model = trainer.get_model(weights=weights, cfg=self.model.yaml)
                        trainer.train()
                    best, last = save_dir / "weights" / "best.pt", save_dir / "weights" / "last.pt"
                    ckpt = best if best.exists() else last
                    metrics = None
                    if self.trainer is not None:
                        metrics = getattr(getattr(trainer, "validator", None), "metrics", None)
                        if metrics is not None:
                            metrics = metrics.results_dict
                    self.metrics[run_name] = metrics or (torch_load(ckpt)["train_metrics"] if ckpt.exists() else None)
                    self.trainers.append(SimpleNamespace(save_dir=save_dir, best=best, last=last))
                except Exception as e:  # one bad dataset should not abort the whole sweep
                    LOGGER.error(f"MultiTrainer: fine-tuning on {data} failed, skipping: {e}")
                    self.metrics[run_name] = None
        finally:
            if base_model:
                base_model.unlink(missing_ok=True)
        if RANK in {-1, 0} and self.trainers:
            self.save_dir.mkdir(parents=True, exist_ok=True)
            self.save_results()  # JSON of per-dataset + mean metrics for programmatic post-processing
            if self.args.get("plots", True):
                self.plot_results()
        return self.metrics

    def save_results(self):
        """Write per-dataset and mean metrics to multitrain_results.json for programmatic post-processing."""
        import json

        results = {run: ({k: float(v) for k, v in m.items()} if m else None) for run, m in self.metrics.items()}
        valid = [m for m in results.values() if m]
        keys = {k for m in valid for k in m}
        self.mean_metrics = {k: sum(m[k] for m in valid if k in m) / sum(k in m for m in valid) for k in keys}
        file = self.save_dir / "multitrain_results.json"
        with open(file, "w", encoding="utf-8") as f:
            json.dump({"results": results, "mean": self.mean_metrics}, f, indent=2)
        LOGGER.info(f"MultiTrainer results saved to {colorstr('bold', file)}")
        return file

    def plot_results(self):
        """Save a cross-dataset bar chart of the per-dataset metric with the mean across all datasets."""
        from ultralytics.cfg import TASK2METRIC
        from ultralytics.utils.plotting import plot_multitrain_results

        key = TASK2METRIC.get(self.args.get("task"))
        scores = {run: float(m.get(key, m.get("fitness", 0.0))) for run, m in self.metrics.items() if m}
        if not scores:
            return None
        fname = plot_multitrain_results(scores, key=key or "fitness", save_dir=self.save_dir)
        LOGGER.info(f"MultiTrainer results saved to {colorstr('bold', fname)}")
        return fname