Skip to content

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=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.

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.

This class combines YOLOETrainer and SegmentationTrainer to provide training functionality specifically for YOLOE segmentation models.

Parameters:

Name Type Description Default
cfg Dict

Configuration dictionary with training parameters.

 DEFAULT_CFG
overrides Dict

Dictionary with parameter overrides.

 None
_callbacks List

List of callback functions for training events.

 None
Source code in ultralytics/models/yolo/yoloe/train_seg.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize the YOLOESegTrainer class.

    This class combines YOLOETrainer and SegmentationTrainer to provide training functionality
    specifically for YOLOE segmentation models.

    Args:
        cfg (Dict): Configuration dictionary with training parameters.
        overrides (Dict, optional): Dictionary with parameter overrides.
        _callbacks (List, optional): List of callback functions for training events.
    """
    if overrides is None:
        overrides = {}
    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 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def get_model(self, cfg=None, weights=None, verbose=True):
    """
    Return YOLOESegModel initialized with specified config and weights.

    Args:
        cfg (dict | str): 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=3,
        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 
68
69
70
71
72
73
74
75
76
77
78
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=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.

Source code in ultralytics/models/yolo/segment/train.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize a SegmentationTrainer object.

    This initializes a trainer for segmentation tasks, extending the detection trainer with segmentation-specific
    functionality. It sets the task to 'segment' and prepares the trainer for training segmentation models.

    Args:
        cfg (dict): Configuration dictionary with default training settings. Defaults to DEFAULT_CFG.
        overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
        _callbacks (list, optional): List of callback functions to be executed during training.

    Examples:
        >>> from ultralytics.models.yolo.segment import SegmentationTrainer
        >>> args = dict(model="yolov8n-seg.pt", data="coco8-seg.yaml", epochs=3)
        >>> trainer = SegmentationTrainer(overrides=args)
        >>> trainer.train()
    """
    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 
 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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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): 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=3,
        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=None, _callbacks=None)

Bases: YOLOETrainerFromScratch, YOLOESegTrainer

Trainer for YOLOE segmentation from scratch.

Source code in ultralytics/models/yolo/yoloe/train.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize the YOLOETrainerFromScratch class.

    This class extends YOLOETrainer to train YOLOE models from scratch. It inherits all functionality from
    the parent class while providing specialized initialization for training without pre-trained weights.

    Args:
        cfg (dict, optional): Configuration dictionary with training parameters. Defaults to DEFAULT_CFG.
        overrides (dict, optional): Dictionary of parameter overrides for configuration.
        _callbacks (list, optional): List of callback functions to be executed during training.

    Examples:
        >>> from ultralytics.models.yoloe.train import YOLOETrainerFromScratch
        >>> trainer = YOLOETrainerFromScratch()
        >>> trainer.train()
    """
    if overrides is None:
        overrides = {}
    super().__init__(cfg, overrides, _callbacks)




ultralytics.models.yolo.yoloe.train_seg.YOLOESegVPTrainer 

YOLOESegVPTrainer(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

Bases: YOLOEVPTrainer, YOLOESegTrainerFromScratch

Trainer for YOLOE segmentation with VP.

Source code in ultralytics/models/yolo/yoloe/train.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize the YOLOETrainerFromScratch class.

    This class extends YOLOETrainer to train YOLOE models from scratch. It inherits all functionality from
    the parent class while providing specialized initialization for training without pre-trained weights.

    Args:
        cfg (dict, optional): Configuration dictionary with training parameters. Defaults to DEFAULT_CFG.
        overrides (dict, optional): Dictionary of parameter overrides for configuration.
        _callbacks (list, optional): List of callback functions to be executed during training.

    Examples:
        >>> from ultralytics.models.yoloe.train import YOLOETrainerFromScratch
        >>> trainer = YOLOETrainerFromScratch()
        >>> trainer.train()
    """
    if overrides is None:
        overrides = {}
    super().__init__(cfg, overrides, _callbacks)



📅 Created 11 days ago ✏️ Updated 11 days ago
RizwanMunawar