Skip to content

Register now for

# Reference for `ultralytics/models/yolo/yoloe/train_seg.py`

Note

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

## `` ultralytics.models.yolo.yoloe.train_seg.YOLOESegTrainer

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

Bases: `YOLOETrainer`, `SegmentationTrainer`

Trainer class for YOLOE segmentation models.

This class combines YOLOETrainer and SegmentationTrainer to provide training functionality specifically for YOLOE segmentation models, enabling both object detection and instance segmentation capabilities.

Attributes:

Name | Type | Description
--- | --- | ---
`cfg` | `dict` | Configuration dictionary with training parameters.
`overrides` | `dict` | Dictionary with parameter overrides.
`_callbacks` | `list` | List of callback functions for training events.

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

    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None):
        """
        Initialize the YOLOE Trainer with specified configurations.

        Args:
            cfg (dict): Configuration dictionary with default training settings from DEFAULT_CFG.
            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
            _callbacks (list, optional): List of callback functions to be applied during training.
        """
        if overrides is None:
            overrides = {}
        assert not overrides.get("compile"), f"Training with 'model={overrides['model']}' requires 'compile=False'"
        overrides["overlap_mask"] = False
        super().__init__(cfg, overrides, _callbacks)

### `` get_model

    get_model(cfg=None, weights=None, verbose=True)

Return YOLOESegModel initialized with specified config and weights.

Parameters:

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

Returns:

Type | Description
--- | ---
`YOLOESegModel` | Initialized YOLOE segmentation model.

Source code in `ultralytics/models/yolo/yoloe/train_seg.py`

    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    def get_model(self, cfg=None, weights=None, verbose=True):
        """
        Return YOLOESegModel initialized with specified config and weights.

        Args:
            cfg (dict | str, optional): Model configuration dictionary or YAML file path.
            weights (str, optional): Path to pretrained weights file.
            verbose (bool): Whether to display model information.

        Returns:
            (YOLOESegModel): Initialized YOLOE segmentation model.
        """
        # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
        # NOTE: Following the official config, nc hard-coded to 80 for now.
        model = YOLOESegModel(
            cfg["yaml_file"] if isinstance(cfg, dict) else cfg,
            ch=self.data["channels"],
            nc=min(self.data["nc"], 80),
            verbose=verbose and RANK == -1,
        )
        if weights:
            model.load(weights)

        return model

### `` get_validator

    get_validator()

Create and return a validator for YOLOE segmentation model evaluation.

Returns:

Type | Description
--- | ---
`YOLOESegValidator` | Validator for YOLOE segmentation models.

Source code in `ultralytics/models/yolo/yoloe/train_seg.py`

    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    def get_validator(self):
        """
        Create and return a validator for YOLOE segmentation model evaluation.

        Returns:
            (YOLOESegValidator): Validator for YOLOE segmentation models.
        """
        self.loss_names = "box", "seg", "cls", "dfl"
        return YOLOESegValidator(
            self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
        )

## `` ultralytics.models.yolo.yoloe.train_seg.YOLOEPESegTrainer

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

Bases: `SegmentationTrainer`

Fine-tune YOLOESeg model in linear probing way.

This trainer specializes in fine-tuning YOLOESeg models using a linear probing approach, which involves freezing most of the model and only training specific layers for efficient adaptation to new tasks.

Attributes:

Name | Type | Description
--- | --- | ---
`data` | `dict` | Dataset configuration containing channels, class names, and number of classes.

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=None, weights=None, verbose=True)

Return YOLOESegModel initialized with specified config and weights for linear probing.

Parameters:

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

Returns:

Type | Description
--- | ---
`YOLOESegModel` | Initialized YOLOE segmentation model configured for linear probing.

Source code in `ultralytics/models/yolo/yoloe/train_seg.py`

     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    def get_model(self, cfg=None, weights=None, verbose=True):
        """
        Return YOLOESegModel initialized with specified config and weights for linear probing.

        Args:
            cfg (dict | str, optional): Model configuration dictionary or YAML file path.
            weights (str, optional): Path to pretrained weights file.
            verbose (bool): Whether to display model information.

        Returns:
            (YOLOESegModel): Initialized YOLOE segmentation model configured for linear probing.
        """
        # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
        # NOTE: Following the official config, nc hard-coded to 80 for now.
        model = YOLOESegModel(
            cfg["yaml_file"] if isinstance(cfg, dict) else cfg,
            ch=self.data["channels"],
            nc=self.data["nc"],
            verbose=verbose and RANK == -1,
        )

        del model.model[-1].savpe

        assert weights is not None, "Pretrained weights must be provided for linear probing."
        if weights:
            model.load(weights)

        model.eval()
        names = list(self.data["names"].values())
        # NOTE: `get_text_pe` related to text model and YOLOEDetect.reprta,
        # it'd get correct results as long as loading proper pretrained weights.
        tpe = model.get_text_pe(names)
        model.set_classes(names, tpe)
        model.model[-1].fuse(model.pe)
        model.model[-1].cv3[0][2] = deepcopy(model.model[-1].cv3[0][2]).requires_grad_(True)
        model.model[-1].cv3[1][2] = deepcopy(model.model[-1].cv3[1][2]).requires_grad_(True)
        model.model[-1].cv3[2][2] = deepcopy(model.model[-1].cv3[2][2]).requires_grad_(True)
        del model.pe
        model.train()

        return model

## `` ultralytics.models.yolo.yoloe.train_seg.YOLOESegTrainerFromScratch

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

Bases: `YOLOETrainerFromScratch`, `YOLOESegTrainer`

Trainer for YOLOE segmentation models trained from scratch without pretrained weights.

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

    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None):
        """
        Initialize the YOLOE Trainer with specified configurations.

        Args:
            cfg (dict): Configuration dictionary with default training settings from DEFAULT_CFG.
            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
            _callbacks (list, optional): List of callback functions to be applied during training.
        """
        if overrides is None:
            overrides = {}
        assert not overrides.get("compile"), f"Training with 'model={overrides['model']}' requires 'compile=False'"
        overrides["overlap_mask"] = False
        super().__init__(cfg, overrides, _callbacks)

## `` ultralytics.models.yolo.yoloe.train_seg.YOLOESegVPTrainer

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

Bases: `YOLOEVPTrainer`, `YOLOESegTrainerFromScratch`

Trainer for YOLOE segmentation models with Vision Prompt (VP) capabilities.

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

    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None):
        """
        Initialize the YOLOE Trainer with specified configurations.

        Args:
            cfg (dict): Configuration dictionary with default training settings from DEFAULT_CFG.
            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
            _callbacks (list, optional): List of callback functions to be applied during training.
        """
        if overrides is None:
            overrides = {}
        assert not overrides.get("compile"), f"Training with 'model={overrides['model']}' requires 'compile=False'"
        overrides["overlap_mask"] = False
        super().__init__(cfg, overrides, _callbacks)

📅 Created 5 months ago ✏️ Updated 5 months ago