Reference for `ultralytics/models/rtdetr/train.py`

Note

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

ultralytics.models.rtdetr.train.RTDETRTrainer

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

Bases: DetectionTrainer

Trainer class for the RT-DETR model developed by Baidu for real-time object detection.

This class extends the DetectionTrainer class for YOLO to adapt to the specific features and architecture of RT-DETR. The model leverages Vision Transformers and has capabilities like IoU-aware query selection and adaptable inference speed.

Attributes:

Name	Type	Description
`loss_names`	`Tuple[str]`	Names of the loss components used for training.
`data`	`dict`	Dataset configuration containing class count and other parameters.
`args`	`dict`	Training arguments and hyperparameters.
`save_dir`	`Path`	Directory to save training results.
`test_loader`	`DataLoader`	DataLoader for validation/testing data.

Notes

F.grid_sample used in RT-DETR does not support the deterministic=True argument.
AMP training can lead to NaN outputs and may produce errors during bipartite graph matching.

Examples:

>>> from ultralytics.models.rtdetr.train import RTDETRTrainer
>>> args = dict(model="rtdetr-l.yaml", data="coco8.yaml", imgsz=640, epochs=3)
>>> trainer = RTDETRTrainer(overrides=args)
>>> trainer.train()

Source code in ultralytics/engine/trainer.py

def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize the BaseTrainer class.

    Args:
        cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CFG.
        overrides (dict, optional): Configuration overrides. Defaults to None.
        _callbacks (list, optional): List of callback functions. Defaults to None.
    """
    self.args = get_cfg(cfg, overrides)
    self.check_resume(overrides)
    self.device = select_device(self.args.device, self.args.batch)
    self.validator = None
    self.metrics = None
    self.plots = {}
    init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)

    # Dirs
    self.save_dir = get_save_dir(self.args)
    self.args.name = self.save_dir.name  # update name for loggers
    self.wdir = self.save_dir / "weights"  # weights dir
    if RANK in {-1, 0}:
        self.wdir.mkdir(parents=True, exist_ok=True)  # make dir
        self.args.save_dir = str(self.save_dir)
        yaml_save(self.save_dir / "args.yaml", vars(self.args))  # save run args
    self.last, self.best = self.wdir / "last.pt", self.wdir / "best.pt"  # checkpoint paths
    self.save_period = self.args.save_period

    self.batch_size = self.args.batch
    self.epochs = self.args.epochs or 100  # in case users accidentally pass epochs=None with timed training
    self.start_epoch = 0
    if RANK == -1:
        print_args(vars(self.args))

    # Device
    if self.device.type in {"cpu", "mps"}:
        self.args.workers = 0  # faster CPU training as time dominated by inference, not dataloading

    # Model and Dataset
    self.model = check_model_file_from_stem(self.args.model)  # add suffix, i.e. yolo11n -> yolo11n.pt
    with torch_distributed_zero_first(LOCAL_RANK):  # avoid auto-downloading dataset multiple times
        self.trainset, self.testset = self.get_dataset()
    self.ema = None

    # Optimization utils init
    self.lf = None
    self.scheduler = None

    # Epoch level metrics
    self.best_fitness = None
    self.fitness = None
    self.loss = None
    self.tloss = None
    self.loss_names = ["Loss"]
    self.csv = self.save_dir / "results.csv"
    self.plot_idx = [0, 1, 2]

    # HUB
    self.hub_session = None

    # Callbacks
    self.callbacks = _callbacks or callbacks.get_default_callbacks()
    if RANK in {-1, 0}:
        callbacks.add_integration_callbacks(self)

build_dataset

build_dataset(img_path, mode='val', batch=None)

Build and return an RT-DETR dataset for training or validation.

Parameters:

Name	Type	Description	Default
`img_path`	`str`	Path to the folder containing images.	required
`mode`	`str`	Dataset mode, either 'train' or 'val'.	`'val'`
`batch`	`int`	Batch size for rectangle training.	`None`

Returns:

Type	Description
`RTDETRDataset`	Dataset object for the specific mode.

Source code in ultralytics/models/rtdetr/train.py

def build_dataset(self, img_path, mode="val", batch=None):
    """
    Build and return an RT-DETR dataset for training or validation.

    Args:
        img_path (str): Path to the folder containing images.
        mode (str): Dataset mode, either 'train' or 'val'.
        batch (int, optional): Batch size for rectangle training.

    Returns:
        (RTDETRDataset): Dataset object for the specific mode.
    """
    return RTDETRDataset(
        img_path=img_path,
        imgsz=self.args.imgsz,
        batch_size=batch,
        augment=mode == "train",
        hyp=self.args,
        rect=False,
        cache=self.args.cache or None,
        single_cls=self.args.single_cls or False,
        prefix=colorstr(f"{mode}: "),
        classes=self.args.classes,
        data=self.data,
        fraction=self.args.fraction if mode == "train" else 1.0,
    )

get_model

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

Initialize and return an RT-DETR model for object detection tasks.

Parameters:

Name	Type	Description	Default
`cfg`	`dict`	Model configuration.	`None`
`weights`	`str`	Path to pre-trained model weights.	`None`
`verbose`	`bool`	Verbose logging if True.	`True`

Returns:

Type	Description
`RTDETRDetectionModel`	Initialized model.

Source code in ultralytics/models/rtdetr/train.py

def get_model(self, cfg=None, weights=None, verbose=True):
    """
    Initialize and return an RT-DETR model for object detection tasks.

    Args:
        cfg (dict, optional): Model configuration.
        weights (str, optional): Path to pre-trained model weights.
        verbose (bool): Verbose logging if True.

    Returns:
        (RTDETRDetectionModel): Initialized model.
    """
    model = RTDETRDetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
    if weights:
        model.load(weights)
    return model

get_validator

get_validator()

Returns a DetectionValidator suitable for RT-DETR model validation.

Source code in ultralytics/models/rtdetr/train.py

def get_validator(self):
    """Returns a DetectionValidator suitable for RT-DETR model validation."""
    self.loss_names = "giou_loss", "cls_loss", "l1_loss"
    return RTDETRValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))

preprocess_batch

preprocess_batch(batch)

Preprocess a batch of images by scaling and converting to float format.

Parameters:

Name	Type	Description	Default
`batch`	`dict`	Dictionary containing a batch of images, bboxes, and labels.	required

Returns:

Type	Description
`dict`	Preprocessed batch with ground truth bounding boxes and classes separated by batch index.

Source code in ultralytics/models/rtdetr/train.py

def preprocess_batch(self, batch):
    """
    Preprocess a batch of images by scaling and converting to float format.

    Args:
        batch (dict): Dictionary containing a batch of images, bboxes, and labels.

    Returns:
        (dict): Preprocessed batch with ground truth bounding boxes and classes separated by batch index.
    """
    batch = super().preprocess_batch(batch)
    bs = len(batch["img"])
    batch_idx = batch["batch_idx"]
    gt_bbox, gt_class = [], []
    for i in range(bs):
        gt_bbox.append(batch["bboxes"][batch_idx == i].to(batch_idx.device))
        gt_class.append(batch["cls"][batch_idx == i].to(device=batch_idx.device, dtype=torch.long))
    return batch

📅 Created 1 year ago ✏️ Updated 7 months ago

Reference for ultralytics/models/rtdetr/train.py

ultralytics.models.rtdetr.train.RTDETRTrainer

build_dataset

get_model

get_validator

preprocess_batch

Reference for `ultralytics/models/rtdetr/train.py`