Skip to content

Register now for

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

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/segment/train.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

## `` ultralytics.models.yolo.segment.train.SegmentationTrainer

    SegmentationTrainer(
        cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None
    )

Bases: `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:

Name | Type | Description
--- | --- | ---
`loss_names` | `tuple[str]` | Names of the loss components used during training.

Examples:

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

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`cfg` | `dict` | Configuration dictionary with default training settings. | `DEFAULT_CFG`
`overrides` | `dict` | Dictionary of parameter overrides for the default configuration. | `None`
`_callbacks` | `list` | List of callback functions to be executed during training. | `None`

Source code in `ultralytics/models/yolo/segment/train.py`

    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=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 (list, optional): List of callback functions to be executed during training.
        """
        if overrides is None:
            overrides = {}
        overrides["task"] = "segment"
        super().__init__(cfg, overrides, _callbacks)

### `` get_model

    get_model(
        cfg: dict | str | None = None,
        weights: str | Path | None = None,
        verbose: bool = True,
    )

Initialize and return a SegmentationModel with specified configuration and weights.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`cfg` | `dict | str` | Model configuration. Can be a dictionary, a path to a YAML file, or None. | `None`
`weights` | `str | Path` | Path to pretrained weights file. | `None`
`verbose` | `bool` | Whether to display model information during initialization. | `True`

Returns:

Type | Description
--- | ---
`SegmentationModel` | Initialized segmentation model with loaded weights if specified.

Examples:

    >>> trainer = SegmentationTrainer()
    >>> model = trainer.get_model(cfg="yolo11n-seg.yaml")
    >>> model = trainer.get_model(weights="yolo11n-seg.pt", verbose=False)

Source code in `ultralytics/models/yolo/segment/train.py`

    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    def get_model(self, cfg: dict | str | None = None, weights: str | Path | 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 (str | Path, optional): Path to pretrained weights file.
            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="yolo11n-seg.yaml")
            >>> model = trainer.get_model(weights="yolo11n-seg.pt", verbose=False)
        """
        model = SegmentationModel(cfg, nc=self.data["nc"], ch=self.data["channels"], verbose=verbose and RANK == -1)
        if weights:
            model.load(weights)

        return model

### `` get_validator

    get_validator()

Return an instance of SegmentationValidator for validation of YOLO model.

Source code in `ultralytics/models/yolo/segment/train.py`

    67
    68
    69
    70
    71
    72
    def get_validator(self):
        """Return an instance of SegmentationValidator for validation of YOLO model."""
        self.loss_names = "box_loss", "seg_loss", "cls_loss", "dfl_loss"
        return yolo.segment.SegmentationValidator(
            self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
        )

📅 Created 1 year ago ✏️ Updated 1 year ago