Reference for ultralytics/models/yolo/segment/train.py#

class SegmentationTrainer(yolo.detect.DetectionTrainer):
    """A class extending the DetectionTrainer class for training based on a segmentation model.

    This trainer specializes in handling segmentation tasks, extending the detection trainer with segmentation-specific
    functionality including model initialization, validation, and visualization.

    Attributes:
        loss_names (tuple[str]): Names of the loss components, derived from the loss dict returned by the criterion.

    Examples:
        >>> from ultralytics.models.yolo.segment import SegmentationTrainer
        >>> args = dict(model="yolo26n-seg.pt", data="coco8-seg.yaml", epochs=3)
        >>> trainer = SegmentationTrainer(overrides=args)
        >>> trainer.train()
    """

    def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks: dict | None = None):
        """Initialize a SegmentationTrainer object.

        Args:
            cfg (dict): Configuration dictionary with default training settings.
            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
            _callbacks (dict, optional): Dictionary of callback functions to be executed during training.
        """
        if overrides is None:
            overrides = {}
        overrides["task"] = "segment"
        super().__init__(cfg, overrides, _callbacks)

def get_model(self, cfg: dict | str | None = None, weights: torch.nn.Module | None = None, verbose: bool = True):
    """Initialize and return a SegmentationModel with specified configuration and weights.

    Args:
        cfg (dict | str, optional): Model configuration. Can be a dictionary, a path to a YAML file, or None.
        weights (torch.nn.Module, optional): Pretrained model whose weights are loaded into the new model.
        verbose (bool): Whether to display model information during initialization.

    Returns:
        (SegmentationModel): Initialized segmentation model with loaded weights if specified.

    Examples:
        >>> trainer = SegmentationTrainer()
        >>> model = trainer.get_model(cfg="yolo26n-seg.yaml", verbose=False)
    """
    model = self.set_model_names_for_load(
        SegmentationModel(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 an instance of SegmentationValidator for validation of YOLO model."""
    return yolo.segment.SegmentationValidator(
        self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
    )