Skip to content

Reference for ultralytics/models/yolo/obb/train.py

Note

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


ultralytics.models.yolo.obb.train.OBBTrainer

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

Bases: DetectionTrainer

A class extending the DetectionTrainer class for training based on an Oriented Bounding Box (OBB) model.

Attributes:

Name Type Description
loss_names Tuple[str]

Names of the loss components used during training.

Methods:

Name Description
get_model

Return OBBModel initialized with specified config and weights.

get_validator

Return an instance of OBBValidator for validation of YOLO model.

Examples:

>>> from ultralytics.models.yolo.obb import OBBTrainer
>>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml", epochs=3)
>>> trainer = OBBTrainer(overrides=args)
>>> trainer.train()

This trainer extends the DetectionTrainer class to specialize in training models that detect oriented bounding boxes. It automatically sets the task to 'obb' in the configuration.

Parameters:

Name Type Description Default
cfg dict

Configuration dictionary for the trainer. Contains training parameters and model configuration.

DEFAULT_CFG
overrides dict

Dictionary of parameter overrides for the configuration. Any values here will take precedence over those in cfg.

None
_callbacks list

List of callback functions to be invoked during training.

None

Examples:

>>> from ultralytics.models.yolo.obb import OBBTrainer
>>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml", epochs=3)
>>> trainer = OBBTrainer(overrides=args)
>>> trainer.train()
Source code in ultralytics/models/yolo/obb/train.py
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize an OBBTrainer object for training Oriented Bounding Box (OBB) models.

    This trainer extends the DetectionTrainer class to specialize in training models that detect oriented
    bounding boxes. It automatically sets the task to 'obb' in the configuration.

    Args:
        cfg (dict, optional): Configuration dictionary for the trainer. Contains training parameters and
            model configuration.
        overrides (dict, optional): Dictionary of parameter overrides for the configuration. Any values here
            will take precedence over those in cfg.
        _callbacks (list, optional): List of callback functions to be invoked during training.

    Examples:
        >>> from ultralytics.models.yolo.obb import OBBTrainer
        >>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml", epochs=3)
        >>> trainer = OBBTrainer(overrides=args)
        >>> trainer.train()
    """
    if overrides is None:
        overrides = {}
    overrides["task"] = "obb"
    super().__init__(cfg, overrides, _callbacks)

get_model

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

Return OBBModel initialized with specified config and weights.

Parameters:

Name Type Description Default
cfg str | dict | None

Model configuration. Can be a path to a YAML config file, a dictionary containing configuration parameters, or None to use default configuration.

None
weights str | Path | None

Path to pretrained weights file. If None, random initialization is used.

None
verbose bool

Whether to display model information during initialization.

True

Returns:

Type Description
OBBModel

Initialized OBBModel with the specified configuration and weights.

Examples:

>>> trainer = OBBTrainer()
>>> model = trainer.get_model(cfg="yolov8n-obb.yaml", weights="yolov8n-obb.pt")
Source code in ultralytics/models/yolo/obb/train.py
def get_model(self, cfg=None, weights=None, verbose=True):
    """
    Return OBBModel initialized with specified config and weights.

    Args:
        cfg (str | dict | None): Model configuration. Can be a path to a YAML config file, a dictionary
            containing configuration parameters, or None to use default configuration.
        weights (str | Path | None): Path to pretrained weights file. If None, random initialization is used.
        verbose (bool): Whether to display model information during initialization.

    Returns:
        (OBBModel): Initialized OBBModel with the specified configuration and weights.

    Examples:
        >>> trainer = OBBTrainer()
        >>> model = trainer.get_model(cfg="yolov8n-obb.yaml", weights="yolov8n-obb.pt")
    """
    model = OBBModel(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 OBBValidator for validation of YOLO model.

Source code in ultralytics/models/yolo/obb/train.py
def get_validator(self):
    """Return an instance of OBBValidator for validation of YOLO model."""
    self.loss_names = "box_loss", "cls_loss", "dfl_loss"
    return yolo.obb.OBBValidator(
        self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
    )



📅 Created 1 year ago ✏️ Updated 7 months ago