Skip to main content
RF-DETR supports training both object detection and segmentation models on custom datasets. You can train using a single train() call or take full control with the Custom Training API built on PyTorch Lightning.

​
Training paths

PathWhen to use
RFDETR.train() (this page)Quickstart, fine-tuning with standard options, Colab notebooks. One call sets up and runs everything.
Custom Training APICustom callbacks, alternative loggers, multi-GPU strategies, integration with external frameworks, or any other customisation of the training loop.
Both paths run the same underlying PyTorch Lightning stack. RFDETR.train() constructs RFDETRModelModule, RFDETRDataModule, and a Trainer internally.

​
Quick start

RF-DETR supports training on datasets in COCO and YOLO formats. The format is automatically detected based on the structure of your dataset directory. 
from rfdetr import RFDETRMedium

model = RFDETRMedium()

model.train(
    dataset_dir="<DATASET_PATH>",
    epochs=100,
    batch_size=4,
    grad_accum_steps=4,
    lr=1e-4,
    output_dir="<OUTPUT_PATH>",
)

​
Effective batch size

The effective batch size determines how many samples the optimizer sees per update: 
effective_batch_size = batch_size × grad_accum_steps × num_gpus
Different GPUs have different VRAM capacities. Adjust batch_size and grad_accum_steps to maintain a total effective batch size of 16. For example:
GPUVRAMbatch_sizegrad_accum_steps
A10040–80 GB161
RTX 409024 GB82
RTX 309024 GB82
T416 GB44
RTX 30708 GB28
To get started quickly with a fine-tuning example, refer to the fine-tuning Colab notebook.

​
Result checkpoints

During training, multiple model checkpoints are saved to the output directory:
FileDescription
checkpoint.pthMost recent checkpoint, saved at the end of the latest epoch. Use this to resume training.
checkpoint_<N>.pthPeriodic checkpoint saved every N epochs (default is every 10).
checkpoint_best_ema.pthBest checkpoint based on validation score using EMA (Exponential Moving Average) weights.
checkpoint_best_regular.pthBest checkpoint based on validation score using raw (non-EMA) model weights.
checkpoint_best_total.pthFinal checkpoint selected for inference. Contains only model weights and is chosen as the better of EMA and non-EMA models.
Training checkpoints (checkpoint.pth, checkpoint_<N>.pth) include model weights, optimizer state, scheduler state, and training metadata — use these to resume training.Evaluation checkpoints (checkpoint_best_ema.pth, checkpoint_best_regular.pth) store only model weights and are used to track the best-performing models.Stripped checkpoint (checkpoint_best_total.pth) contains only the final model weights and is optimized for inference and deployment.

​
Load and run a fine-tuned model

After training completes, load your checkpoint by passing its path to pretrain_weights: 
from rfdetr import RFDETRMedium

model = RFDETRMedium(pretrain_weights="<CHECKPOINT_PATH>")

detections = model.predict("<IMAGE_PATH>")
Use pretrain_weights="checkpoint_best_total.pth" when initializing a model to start fresh training from those weights. Use resume="checkpoint.pth" to continue training with full optimizer state.

​
Next steps

Dataset formats

Prepare your dataset in COCO or YOLO format for RF-DETR training.

Training parameters

Complete reference for all RF-DETR training configuration options.

Advanced training

Resume training, early stopping, multi-GPU, and memory optimization.

Training loggers

Track experiments with TensorBoard, W&B, MLflow, and ClearML.

Build docs developers (and LLMs) love

Get started for freeTalk to us