Skip to content

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

Improvements

This page is sourced from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/yoloe/train_seg.py. Have an improvement or example to add? Open a Pull Request — thank you! 🙏


Summary

class ultralytics.models.yolo.yoloe.train_seg.YOLOESegTrainer

YOLOESegTrainer()

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

NameTypeDescription
cfgdictConfiguration dictionary with training parameters.
overridesdictDictionary with parameter overrides.
_callbackslistList of callback functions for training events.

Methods

NameDescription
get_modelReturn YOLOESegModel initialized with specified config and weights.
get_validatorCreate and return a validator for YOLOE segmentation model evaluation.
Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):


method ultralytics.models.yolo.yoloe.train_seg.YOLOESegTrainer.get_model

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

Return YOLOESegModel initialized with specified config and weights.

Args

NameTypeDescriptionDefault
cfgdict | str, optionalModel configuration dictionary or YAML file path.None
weightsstr, optionalPath to pretrained weights file.None
verboseboolWhether to display model information.True

Returns

TypeDescription
YOLOESegModelInitialized YOLOE segmentation model.
Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
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


method ultralytics.models.yolo.yoloe.train_seg.YOLOESegTrainer.get_validator

def get_validator(self)

Create and return a validator for YOLOE segmentation model evaluation.

Returns

TypeDescription
YOLOESegValidatorValidator for YOLOE segmentation models.
Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
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
    )




class ultralytics.models.yolo.yoloe.train_seg.YOLOEPESegTrainer

YOLOEPESegTrainer()

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

NameTypeDescription
datadictDataset configuration containing channels, class names, and number of classes.

Methods

NameDescription
get_modelReturn YOLOESegModel initialized with specified config and weights for linear probing.
Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
class YOLOEPESegTrainer(SegmentationTrainer):


method ultralytics.models.yolo.yoloe.train_seg.YOLOEPESegTrainer.get_model

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

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

Args

NameTypeDescriptionDefault
cfgdict | str, optionalModel configuration dictionary or YAML file path.None
weightsstr, optionalPath to pretrained weights file.None
verboseboolWhether to display model information.True

Returns

TypeDescription
YOLOESegModelInitialized YOLOE segmentation model configured for linear probing.
Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
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




class ultralytics.models.yolo.yoloe.train_seg.YOLOESegTrainerFromScratch

YOLOESegTrainerFromScratch()

Bases: YOLOETrainerFromScratch, YOLOESegTrainer

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

Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
class YOLOESegTrainerFromScratch(YOLOETrainerFromScratch, YOLOESegTrainer):




class ultralytics.models.yolo.yoloe.train_seg.YOLOESegVPTrainer

YOLOESegVPTrainer()

Bases: YOLOEVPTrainer, YOLOESegTrainerFromScratch

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

Source code in ultralytics/models/yolo/yoloe/train_seg.pyView on GitHub
class YOLOESegVPTrainer(YOLOEVPTrainer, YOLOESegTrainerFromScratch):





📅 Created 8 months ago ✏️ Updated 2 days ago
glenn-jocherRizwanMunawar