Skip to content

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=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()
Source code in ultralytics/models/yolo/segment/train.py 
28
29
30
31
32
33
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """Initialize a SegmentationTrainer object with given arguments."""
    if overrides is None:
        overrides = {}
    overrides["task"] = "segment"
    super().__init__(cfg, overrides, _callbacks)

get_model 

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

Return SegmentationModel initialized with specified config and weights.

Source code in ultralytics/models/yolo/segment/train.py 
35
36
37
38
39
40
41
def get_model(self, cfg=None, weights=None, verbose=True):
    """Return SegmentationModel initialized with specified config and weights."""
    model = SegmentationModel(cfg, ch=3, nc=self.data["nc"], 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 
43
44
45
46
47
48
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
    )

plot_metrics 

plot_metrics()

Plots training/val metrics.

Source code in ultralytics/models/yolo/segment/train.py 
63
64
65
def plot_metrics(self):
    """Plots training/val metrics."""
    plot_results(file=self.csv, segment=True, on_plot=self.on_plot)  # save results.png

plot_training_samples 

plot_training_samples(batch, ni)

Creates a plot of training sample images with labels and box coordinates.

Source code in ultralytics/models/yolo/segment/train.py 
50
51
52
53
54
55
56
57
58
59
60
61
def plot_training_samples(self, batch, ni):
    """Creates a plot of training sample images with labels and box coordinates."""
    plot_images(
        batch["img"],
        batch["batch_idx"],
        batch["cls"].squeeze(-1),
        batch["bboxes"],
        masks=batch["masks"],
        paths=batch["im_file"],
        fname=self.save_dir / f"train_batch{ni}.jpg",
        on_plot=self.on_plot,
    )



📅 Created 1 year ago ✏️ Updated 6 months ago
jk4e glenn-jocher Burhan-Q