# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license import math import random from copy import copy import numpy as np import torch.nn as nn from ultralytics.data import build_dataloader, build_yolo_dataset from ultralytics.engine.trainer import BaseTrainer from ultralytics.models import yolo from ultralytics.nn.tasks import DetectionModel from ultralytics.utils import LOGGER, RANK from ultralytics.utils.plotting import plot_images, plot_labels, plot_results from ultralytics.utils.torch_utils import de_parallel, torch_distributed_zero_first class DetectionTrainer(BaseTrainer): """ A class extending the BaseTrainer class for training based on a detection model. This trainer specializes in object detection tasks, handling the specific requirements for training YOLO models for object detection. Attributes: model (DetectionModel): The YOLO detection model being trained. data (dict): Dictionary containing dataset information including class names and number of classes. loss_names (Tuple[str]): Names of the loss components used in training (box_loss, cls_loss, dfl_loss). Methods: build_dataset: Build YOLO dataset for training or validation. get_dataloader: Construct and return dataloader for the specified mode. preprocess_batch: Preprocess a batch of images by scaling and converting to float. set_model_attributes: Set model attributes based on dataset information. get_model: Return a YOLO detection model. get_validator: Return a validator for model evaluation. label_loss_items: Return a loss dictionary with labeled training loss items. progress_string: Return a formatted string of training progress. plot_training_samples: Plot training samples with their annotations. plot_metrics: Plot metrics from a CSV file. plot_training_labels: Create a labeled training plot of the YOLO model. auto_batch: Calculate optimal batch size based on model memory requirements. Examples: >>> from ultralytics.models.yolo.detect import DetectionTrainer >>> args = dict(model="yolo11n.pt", data="coco8.yaml", epochs=3) >>> trainer = DetectionTrainer(overrides=args) >>> trainer.train() """ def build_dataset(self, img_path, mode="train", batch=None): """ Build YOLO Dataset for training or validation. Args: img_path (str): Path to the folder containing images. mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode. batch (int, optional): Size of batches, this is for `rect`. Returns: (Dataset): YOLO dataset object configured for the specified mode. """ gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32) return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs) def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"): """ Construct and return dataloader for the specified mode. Args: dataset_path (str): Path to the dataset. batch_size (int): Number of images per batch. rank (int): Process rank for distributed training. mode (str): 'train' for training dataloader, 'val' for validation dataloader. Returns: (DataLoader): PyTorch dataloader object. """ assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}." with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP dataset = self.build_dataset(dataset_path, mode, batch_size) shuffle = mode == "train" if getattr(dataset, "rect", False) and shuffle: LOGGER.warning("'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False") shuffle = False workers = self.args.workers if mode == "train" else self.args.workers * 2 return build_dataloader(dataset, batch_size, workers, shuffle, rank) # return dataloader def preprocess_batch(self, batch): """ Preprocess a batch of images by scaling and converting to float. Args: batch (dict): Dictionary containing batch data with 'img' tensor. Returns: (dict): Preprocessed batch with normalized images. """ batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255 if self.args.multi_scale: imgs = batch["img"] sz = ( random.randrange(int(self.args.imgsz * 0.5), int(self.args.imgsz * 1.5 + self.stride)) // self.stride * self.stride ) # size sf = sz / max(imgs.shape[2:]) # scale factor if sf != 1: ns = [ math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:] ] # new shape (stretched to gs-multiple) imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) batch["img"] = imgs return batch def set_model_attributes(self): """Set model attributes based on dataset information.""" # Nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps) # self.args.box *= 3 / nl # scale to layers # self.args.cls *= self.data["nc"] / 80 * 3 / nl # scale to classes and layers # self.args.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl # scale to image size and layers self.model.nc = self.data["nc"] # attach number of classes to model self.model.names = self.data["names"] # attach class names to model self.model.args = self.args # attach hyperparameters to model # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc def get_model(self, cfg=None, weights=None, verbose=True): """ Return a YOLO detection model. Args: cfg (str, optional): Path to model configuration file. weights (str, optional): Path to model weights. verbose (bool): Whether to display model information. Returns: (DetectionModel): YOLO detection model. """ model = DetectionModel(cfg, nc=self.data["nc"], ch=self.data["channels"], verbose=verbose and RANK == -1) if weights: model.load(weights) return model def get_validator(self): """Return a DetectionValidator for YOLO model validation.""" self.loss_names = "box_loss", "cls_loss", "dfl_loss" return yolo.detect.DetectionValidator( self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks ) def label_loss_items(self, loss_items=None, prefix="train"): """ Return a loss dict with labeled training loss items tensor. Args: loss_items (List[float], optional): List of loss values. prefix (str): Prefix for keys in the returned dictionary. Returns: (Dict | List): Dictionary of labeled loss items if loss_items is provided, otherwise list of keys. """ keys = [f"{prefix}/{x}" for x in self.loss_names] if loss_items is not None: loss_items = [round(float(x), 5) for x in loss_items] # convert tensors to 5 decimal place floats return dict(zip(keys, loss_items)) else: return keys def progress_string(self): """Return a formatted string of training progress with epoch, GPU memory, loss, instances and size.""" return ("\n" + "%11s" * (4 + len(self.loss_names))) % ( "Epoch", "GPU_mem", *self.loss_names, "Instances", "Size", ) def plot_training_samples(self, batch, ni): """ Plot training samples with their annotations. Args: batch (dict): Dictionary containing batch data. ni (int): Number of iterations. """ plot_images( images=batch["img"], batch_idx=batch["batch_idx"], cls=batch["cls"].squeeze(-1), bboxes=batch["bboxes"], paths=batch["im_file"], fname=self.save_dir / f"train_batch{ni}.jpg", on_plot=self.on_plot, ) def plot_metrics(self): """Plot metrics from a CSV file.""" plot_results(file=self.csv, on_plot=self.on_plot) # save results.png def plot_training_labels(self): """Create a labeled training plot of the YOLO model.""" boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0) cls = np.concatenate([lb["cls"] for lb in self.train_loader.dataset.labels], 0) plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot) def auto_batch(self): """ Get optimal batch size by calculating memory occupation of model. Returns: (int): Optimal batch size. """ train_dataset = self.build_dataset(self.trainset, mode="train", batch=16) max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4 # 4 for mosaic augmentation return super().auto_batch(max_num_obj)