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: Optional[Dict] = 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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict] = None, _callbacks=None):
    """
    Initialize the YOLOE Trainer with specified configurations.

    This method sets up the YOLOE trainer with the provided configuration and overrides, initializing
    the training environment, model, and callbacks for YOLOE object detection training.

    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 = {}
    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: Optional[Dict] = 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
43
44
45
46
47
48
49
50
51
def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict] = 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.
        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="yolo11n-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 
 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: Optional[Dict] = 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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict] = None, _callbacks=None):
    """
    Initialize the YOLOE Trainer with specified configurations.

    This method sets up the YOLOE trainer with the provided configuration and overrides, initializing
    the training environment, model, and callbacks for YOLOE object detection training.

    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 = {}
    overrides["overlap_mask"] = False
    super().__init__(cfg, overrides, _callbacks)




ultralytics.models.yolo.yoloe.train_seg.YOLOESegVPTrainer 

YOLOESegVPTrainer(
    cfg=DEFAULT_CFG, overrides: Optional[Dict] = 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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict] = None, _callbacks=None):
    """
    Initialize the YOLOE Trainer with specified configurations.

    This method sets up the YOLOE trainer with the provided configuration and overrides, initializing
    the training environment, model, and callbacks for YOLOE object detection training.

    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 = {}
    overrides["overlap_mask"] = False
    super().__init__(cfg, overrides, _callbacks)





📅 Created 3 months ago ✏️ Updated 3 months ago
RizwanMunawar