Train an RFDETR Model

SUMMARY

Train and fine-tune RF-DETR object-detection and instance-segmentation models with Iris. Select an RFDETRVariant, configure RFDETRConfig, prepare COCO datasets, and run or resume training with the shared Trainer.

Python Version Support

Iris training supports Python 3.11 and 3.12 only. Create an environment with one of these versions before installing dependencies.

Requirements

Requirement	Value
Python	3.11 or 3.12
Platform	Linux, macOS, or Windows
Dataset	COCO detection or instance-segmentation annotations
GPU	CUDA GPU recommended for practical training times

Quick Installation

Install PyTorch for your CUDA version, then install the RF-DETR dependencies:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
pip install "telekinesis-ai[iris-rfdetr]"

PyTorch Installation

For CPU, macOS, or another CUDA version, select the appropriate command in the PyTorch installation guide.

Training Workflow

The example below configures an RF-DETR segmentation model, creates training and validation datasets, records TensorBoard metrics, and optionally resumes from a checkpoint.

Create train.py

"""Train an RF-DETR model on a COCO dataset."""

import argparse
from pathlib import Path

import torch
from loguru import logger

from telekinesis.iris import dataset, logger as iris_logger, models, trainer


def train_custom_rfdetr_model_example(args: argparse.Namespace) -> None:
    """Configure and train an RF-DETR model on a COCO dataset."""
    seed = 42

    # ===================== Configure Model ================================
    # Detection: nano, small, medium, base, large
    # Segmentation: seg-nano, seg-small, seg-medium, seg-large,
    #               seg-xlarge, seg-2xlarge
    model_config = models.RFDETRConfig(
        variant=models.RFDETRVariant.SEG_MEDIUM,
        num_classes=5,
        pretrained=True,
        resolution=None,
        validation_images_to_log=4,
    )
    model = models.RFDETR(config=model_config)

    # ===================== Prepare Dataset ================================
    training_dataset, validation_dataset = dataset.prepare_coco_datasets(
        dataset_dir=args.dataset_dir,
        validation_split=0.3,
        seed=seed,
        train_transforms=model.train_transforms,
        val_transforms=model.val_transforms,
        include_masks=model.requires_masks,
    )

    # ===================== Train Model ====================================
    output_dir = args.output_dir
    device = "cuda" if torch.cuda.is_available() else "cpu"

    logger.info(
        f"Training RF-DETR {model_config.variant.value} on "
        f"{len(training_dataset)} images at "
        f"{model.model_config.resolution}x{model.model_config.resolution} "
        f"using {device}."
    )

    metric_logger = iris_logger.TensorBoardMetricLogger(
        output_dir / "tensorboard"
    )
    # Tune these values for your dataset size and available GPU memory.
    model_trainer = trainer.Trainer(
        model=model,
        training_dataset=training_dataset,
        validation_dataset=validation_dataset,
        output_dir=output_dir,
        epochs=400,
        batch_size=8,
        learning_rate=1e-4,
        weight_decay=1e-4,
        device=device,
        num_workers=0,
        mixed_precision=True,
        evaluation_interval=1,
        checkpoint_interval=1,
        seed=seed,
        metric_logger=metric_logger,
    )

    try:
        if args.resume_from is None:
            model_trainer.train()
        else:
            model_trainer.resume(args.resume_from)
    except KeyboardInterrupt:
        logger.info("Training interrupted; existing checkpoints were preserved.")
    except Exception as error:
        logger.error(f"An error occurred during training: {error}")

    tensorboard_dir = (output_dir / "tensorboard").resolve()
    logger.info(f'TensorBoard: tensorboard --logdir "{tensorboard_dir}"')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        required=True,
        help="Path to the COCO dataset directory.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help="Directory in which to save training outputs.",
    )
    parser.add_argument(
        "--resume-from",
        type=Path,
        help="Optional checkpoint from which to resume training.",
    )
    train_custom_rfdetr_model_example(parser.parse_args())

Launch Training

Start a new training run:

python train.py \
  --dataset-dir path/to/dataset \
  --output-dir results/rfdetr

Resume a run from a saved checkpoint:

python train.py \
  --dataset-dir path/to/dataset \
  --output-dir results/rfdetr \
  --resume-from results/rfdetr/latest.pt

Training Guidance

See Best Practices: Trainer for guidance on tuning batch size, learning rate, epochs, workers, evaluation, and checkpoints.

Class Count and Masks

Set num_classes to the classification-head size required by the dataset category IDs. Segmentation variants also require a valid polygon or RLE segmentation for every annotated object.

Run the Example

A runnable version is available in the Telekinesis examples repository:

cd telekinesis-examples
python examples/iris/train_custom_rfdetr_model.py \
  --dataset-dir path/to/dataset \
  --output-dir results/rfdetr

Add --resume-from results/rfdetr/latest.pt to continue a previous run.

Checkpoints and Metrics

The training script saves the following outputs to the configured output_dir:

Output	Description
latest.pt	Updated after every completed epoch
best.pt	Updated when the monitored validation or training loss improves
epoch_NNNN.pt	Written according to checkpoint_interval
TensorBoard logs	Training loss, learning rate, validation metrics, and configured previews

Validation reports COCO box metrics. Segmentation variants additionally report mask metrics.

Model Reference

Model Variants

Pass an RFDETRVariant enum member—or its string value—to RFDETRConfig.variant.

Task	Enum member	Value
Object detection	RFDETRVariant.NANO	"nano"
Object detection	RFDETRVariant.SMALL	"small"
Object detection	RFDETRVariant.MEDIUM	"medium"
Object detection	RFDETRVariant.BASE	"base"
Object detection	RFDETRVariant.LARGE	"large"
Instance segmentation	RFDETRVariant.SEG_NANO	"seg-nano"
Instance segmentation	RFDETRVariant.SEG_SMALL	"seg-small"
Instance segmentation	RFDETRVariant.SEG_MEDIUM	"seg-medium"
Instance segmentation	RFDETRVariant.SEG_LARGE	"seg-large"

Start with a smaller variant while validating a new dataset and pipeline, then increase model size if accuracy requires it.

RFDETRConfig

RFDETRConfig combines model construction and training-adapter settings in one immutable configuration.

Option	Default	Description
variant	RFDETRVariant.SEG_MEDIUM	RF-DETR detection or segmentation variant
num_classes	1	Positive number of foreground class slots
pretrained	True	Loads pretrained model weights
resolution	None	Square input resolution; uses the variant default when omitted
num_queries	None	Overrides the variant's number of object queries
gradient_checkpointing	None	Overrides activation checkpointing
compile	None	Overrides torch.compile behavior
freeze_encoder	None	Overrides encoder freezing
validation_images_to_log	4	Maximum validation previews logged per evaluation
image_mean	(0.485, 0.456, 0.406)	RGB normalization mean; None disables mean normalization
image_std	(0.229, 0.224, 0.225)	Positive RGB normalization standard deviation
box_format	"xyxy"	Adapter input box format: "xyxy" or "cxcywh"
normalize_boxes	True	Normalizes bounding-box coordinates for training
background_class_id	0	Background category ID; None disables background filtering

Custom Resolution

Custom resolutions must satisfy the selected variant's patch/window divisibility constraints. Prefer the variant default unless the input scale or available memory requires an override.

RFDETR Properties

Property	Type	Description
config	RFDETRConfig	Public unified configuration
model_config	ModelConfig	Resolved variant-specific model configuration
adapter_config	RFDETRConfig	Configuration used by the training adapter
model	nn.Module	Underlying PyTorch model
adapter	RFDETRTrainingAdapter	Cached RF-DETR training adapter
name	str	Selected variant string
train_transforms	ResizeTransforms	Training resize transforms
val_transforms	ResizeTransforms	Validation resize transforms
requires_masks	bool	True when the selected variant has a segmentation head

Trainer Reference

See Configuration: Trainer Reference for all Trainer constructor options, defaults, and methods.

RFDETR Best Practices

Categories and Annotations

• Keep category IDs consistent between training and validation datasets.
• Set num_classes to the classification-head size required by the category IDs. With standard COCO IDs starting at 1, this is typically the greatest category ID.
• Keep background_class_id=0 for standard COCO datasets without a foreground category at ID 0. Set it to None only when no class slot should be filtered as background.
• Use a segmentation variant only when every annotated object has a valid polygon or RLE mask.
• Inspect boxes and masks visually before starting a long training run.

Variant Selection

• Start with NANO or SEG_NANO while validating a new dataset and pipeline.
• Increase model size only when validation metrics show that the smaller variant lacks capacity.
• Use a detection variant when masks are not required; segmentation variants add memory and compute overhead.
• Compare variants with the same dataset split and Trainer seed.

config = RFDETRConfig(
    variant=RFDETRVariant.SEG_NANO,
    num_classes=5,
    pretrained=True,
)

Pretrained Weights and Encoder Freezing

• Keep pretrained=True for normal fine-tuning on a custom dataset.
• Use freeze_encoder=True when compute or data is limited and only the prediction layers need adaptation.
• Unfreeze the encoder when the deployment images differ substantially from the pretraining domain.
• Validate the effect of encoder freezing rather than assuming it will improve generalization.

Resolution and Object Queries

• Prefer the variant's default resolution initially.
• Increase resolution when small-object detail is being lost; expect higher memory use and slower training.
• Reduce Trainer batch_size before lowering resolution when GPU memory is limited.
• Custom resolutions must be divisible by the selected variant's patch-size and window configuration.
• Increase num_queries only when images routinely contain more objects than the variant can represent. Extra queries increase memory and compute.

config = RFDETRConfig(
    variant=RFDETRVariant.SEG_MEDIUM,
    num_classes=5,
    resolution=None,
    num_queries=None,
)

Memory and Performance

• Enable gradient_checkpointing when activation memory limits training. It reduces memory use at the cost of additional computation.
• Treat compile=True as an optimization to benchmark on the target environment; compilation adds startup cost and may not support every platform equally.
• Reduce batch_size first when addressing out-of-memory errors.
• Keep mixed precision enabled on CUDA unless numerical instability requires full precision.
• Change one performance setting at a time and compare both throughput and validation metrics.

Preprocessing and Boxes

• Keep the default ImageNet image_mean and image_std when using pretrained weights unless the input pipeline has a documented alternative normalization.
• COCODataset supplies boxes in xyxy, matching the default box_format="xyxy".
• Keep normalize_boxes=True for the standard RF-DETR training adapter.
• Change box_format, normalization, or RGB statistics only when using a custom dataset pipeline that requires it.

Validation and Resume

• Use validation_images_to_log to inspect qualitative predictions alongside COCO metrics.
• Select deployment candidates using validation results rather than training loss alone.
• Use best.pt for the best monitored model and latest.pt to continue interrupted training.
• Keep the variant, class count, resolution, query count, and preprocessing configuration compatible when resuming.

Next Steps

Export a Model

Export an RFDETR checkpoint as an ONNX model.

Export a model →

Train SAM3-LoRA

Fine-tune prompted instance segmentation with compact LoRA weights.

Train SAM3-LoRA →