---
description: Train RF-DETR detection and segmentation models on custom datasets. Supports COCO and YOLO formats with one-line Python API and PyTorch Lightning.
---

# Train an RF-DETR Model

!!! tip "Key Takeaways"

    - Train detection, segmentation, or keypoint preview models with a single `model.train(dataset_dir=...)` call
    - Detection and segmentation support COCO JSON and YOLO dataset formats with automatic detection
    - Keypoint preview training supports COCO keypoint JSON and Ultralytics YOLO pose datasets
    - Fine-tune from COCO-pretrained checkpoints (Nano to 2XLarge) for fastest convergence
    - Built on PyTorch Lightning — use the high-level API or access PTL primitives directly for full control
    - EMA weights, early stopping, and best-model checkpointing are included by default

You can train RF-DETR object detection and segmentation models on a custom dataset using the `rfdetr` Python package, or in the cloud using Roboflow.

This guide describes how to train both an object detection and segmentation RF-DETR model.

## Training paths

RF-DETR provides two training paths:

| Path                                        | When to use                                                                                                                                         |
| ------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------- |
| **`RFDETR.train()`** (this page)            | Quickstart, fine-tuning with standard options, Colab notebooks. One call sets up and runs everything.                                               |
| **[Custom Training API](customization.md)** | Custom 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; the Lightning API page shows how to do the same thing explicitly so you can modify each component.

## Quick Start

!!! info "Training requires the `train` extra"

    Training dependencies are not included in the base install. Install them with:

    ```bash
    pip install "rfdetr[train]"
    ```

    For experiment tracking, also add `pip install "rfdetr[train,loggers]"`.

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

=== "Object Detection"

    ```python
    from rfdetr import RFDETRMedium

    model = RFDETRMedium()

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

=== "Image Segmentation"

    ```python
    from rfdetr import RFDETRSegMedium

    model = RFDETRSegMedium()

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

=== "Keypoint Preview"

    ```python
    from rfdetr import RFDETRKeypointPreview

    model = RFDETRKeypointPreview()

    model.train(
        dataset_dir="<KEYPOINT_DATASET_PATH>",
        epochs=50,
        batch_size="auto",
        lr=1e-5,
        output_dir="<OUTPUT_PATH>",
    )
    ```

Different models, tasks, resolutions, and GPUs have different memory requirements. On CUDA, the portable starting point is `batch_size="auto"`: the probe sizes the physical batch for the current training setup and recommends accumulation toward its configured target. CPU and MPS training require a concrete integer batch size; if hardware memory limits that physical batch, increase `grad_accum_steps` to recover a nominal effective-batch target. The product is not an optimization-equivalence guarantee: accumulation changes forward/backward microbatch cadence, and a larger physical batch can have different throughput and optimization behavior.

Each model class downloads its COCO-pretrained checkpoint automatically when instantiated. To get started quickly with training an object detection model, please refer to our fine-tuning Google Colab [notebook](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-rf-detr-on-detection-dataset.ipynb).

## Keypoint preview custom datasets

The pretrained keypoint preview checkpoint predicts 17 COCO person keypoints. Fine-tuned keypoint preview models can use the keypoint schema from your own COCO or YOLO pose dataset, so the output keypoint count is not limited to 17.

Use COCO keypoint JSON or Ultralytics YOLO pose labels for custom keypoint training. Roboflow COCO exports are supported when split annotations are named `train/_annotations.coco.json`, `valid/_annotations.coco.json`, and optionally `test/_annotations.coco.json`. YOLO pose datasets use the existing RF-DETR YOLO directory layout with `data.yaml`, `train/images`, `train/labels`, `valid/images`, and `valid/labels`.

The keypoint fine-tuning demo infers the class names and keypoint schema from the training annotation file, then passes those values into `RFDETRKeypointPreview` and `model.train()`:

```python
from pathlib import Path

from rfdetr import RFDETRKeypointPreview
from rfdetr.datasets._keypoint_schema import infer_coco_keypoint_schema

DATASET_DIR = Path("/path/to/coco-keypoint-dataset")
schema = infer_coco_keypoint_schema(DATASET_DIR / "train" / "_annotations.coco.json")

model = RFDETRKeypointPreview(
    num_classes=len(schema.class_names),
    num_keypoints_per_class=schema.num_keypoints_per_class,
)

model.train(
    dataset_file="roboflow",
    dataset_dir=str(DATASET_DIR),
    class_names=schema.class_names,
    keypoint_oks_sigmas=schema.keypoint_oks_sigmas,
    epochs=50,
    batch_size=8,
    grad_accum_steps=2,
    lr=2e-5,
    lr_encoder=2e-5,
    output_dir="output/keypoint_custom",
    use_ema=False,
    run_test=False,
)
```

Set `keypoint_flip_pairs` if horizontal flips should swap left/right keypoints for your schema.

For YOLO pose datasets, use `infer_yolo_keypoint_schema(DATASET_DIR / "data.yaml")` instead. RF-DETR also infers YOLO pose schema automatically during `model.train()` when `data.yaml` declares `kpt_shape`.

## Dataset Format

RF-DETR **automatically detects** whether your dataset is in COCO or YOLO format. Simply pass your dataset directory to the `train()` method and the appropriate data loader will be used.

| Format   | Detection Method                                               | Learn More                                          |
| -------- | -------------------------------------------------------------- | --------------------------------------------------- |
| **COCO** | Looks for `train/_annotations.coco.json`                       | [COCO Format Guide](dataset-formats.md#coco-format) |
| **YOLO** | Looks for `data.yaml` or `data.yml` + resolved training images | [YOLO Format Guide](dataset-formats.md#yolo-format) |

For keypoint preview training, use COCO keypoint JSON or YOLO pose labels. YOLO pose datasets must declare `kpt_shape` in `data.yaml`; detection-only YOLO datasets still fail clearly in keypoint mode instead of being treated as pose labels.

[Roboflow](https://roboflow.com/annotate) allows you to create object detection datasets from scratch and export them in either COCO JSON or YOLO format for training. You can also explore [Roboflow Universe](https://universe.roboflow.com/) to find pre-labeled datasets for a range of use cases.

→ **[Learn more about dataset formats](dataset-formats.md)**

## Training Configuration

RF-DETR provides many configuration options to customize your training run. See the complete reference for all available parameters.

→ **[View all training parameters](training-parameters.md)**

## Advanced Topics

- [Resume training](advanced.md#resume-training) from a checkpoint
- [Early stopping](advanced.md#early-stopping) to prevent overfitting
- [Multi-GPU training](advanced.md#multi-gpu-training) with PyTorch Lightning DDP
- [Default and custom augmentations](augmentations.md) - Torchvision defaults plus optional Albumentations (CPU) or Kornia (GPU) configs
- [Memory optimization](advanced.md#memory-optimization) with gradient checkpointing

→ **[Learn more about advanced training](advanced.md)**

## Custom Training API

RF-DETR's training stack is built on PyTorch Lightning. The `RFDETR.train()` call above constructs and runs PTL primitives internally. Use them directly when you need custom callbacks, non-default loggers, multi-GPU strategies, or full control over the training loop.

→ **[Custom Training API guide](customization.md)**

## Training Loggers

Track your experiments with popular logging platforms:

- [TensorBoard](loggers.md#tensorboard) for local visualization
- [Weights and Biases](loggers.md#weights-and-biases) for cloud-based tracking
- [ClearML](loggers.md#clearml) workaround for SDK auto-binding
- [MLflow](loggers.md#mlflow) for experiment lifecycle management

→ **[Learn more about training loggers](loggers.md)**

## Result Checkpoints

During training, multiple model checkpoints are saved to the output directory:

- `last.ckpt` – the most recent full checkpoint, saved at the end of the latest epoch.

- `checkpoint_<epoch>.ckpt` – periodic full checkpoints saved every N epochs (default is every 10).

    Both `.ckpt` files also store the training config, the model class name and the model config, the same fields as the `.pth` files below, so `RFDETR.from_checkpoint("output/last.ckpt")` rebuilds the model with that epoch's non-EMA weights and the dataset's class names.

- `checkpoint_best_ema.pth` – best checkpoint based on validation score, using the EMA (Exponential Moving Average) weights. EMA weights are a smoothed version of the model's parameters across training steps, often yielding better generalization.

- `checkpoint_best_regular.pth` – best checkpoint based on validation score, using the raw (non-EMA) model weights.

- `checkpoint_best_total.pth` – final checkpoint selected for inference and benchmarking. It contains model weights, the model config (so `RFDETR.from_checkpoint()` rebuilds the trained resolution and architecture), epoch/PTL metadata, and callback state when available, but no optimizer or scheduler state. It is chosen as the better of the EMA and non-EMA models based on validation performance.

- `last_ema.pth` – final EMA weights, written at the end of training when EMA is enabled (the default). Mirrors `last.ckpt` for the EMA model.

For detection models, the validation score is box mAP (`val/mAP_50_95`). For segmentation models, it is mask mAP (`val/segm_mAP_50_95`), not box mAP. For keypoint preview models, best-checkpoint selection uses COCO keypoint AP (`val/keypoint_map_50_95`) and checkpoints persist the model keypoint schema so `RFDETR.from_checkpoint()` can reconstruct the same label/keypoint slots. `TrainConfig.best_model_metric="mar"` selects on mAR instead — unlike mAP, detection and segmentation share the same unprefixed `val/mAR` key; keypoint models still use their own `val/keypoint_mAR`.

The output directory also holds `training_config.json`, a record of how the run was configured: the full `TrainConfig` and `ModelConfig`, the model config class name, and the dataset's class names. It is written when training starts, so a run you interrupt with Ctrl-C still leaves one behind, and rewritten when training finishes — at which point `class_names` is the list you passed as `class_names`, or the dataset's own labels when you passed none. Because the start-of-run copy is written before the first batch, the file's presence does not by itself mean a run reached training, and it replaces any copy an earlier run left in the same `output_dir`.

??? note "Checkpoint file sizes"

    Checkpoint sizes vary based on what they contain:

    - **Training checkpoints** (e.g. `last.ckpt`, `checkpoint_<epoch>.ckpt`) include model weights, optimizer state, scheduler state, and training metadata. Use these to resume training.

    - **Lightweight best checkpoints** (e.g. `checkpoint_best_ema.pth`, `checkpoint_best_regular.pth`, `last_ema.pth`) store model weights, epoch/PTL metadata, and callback state when available, but intentionally omit optimizer and scheduler state. These may come from different epochs depending on which version achieved the highest validation score.

    - **Lightweight total checkpoint** (e.g. `checkpoint_best_total.pth`) keeps the same lightweight resume metadata while selecting the final best model for inference and deployment.

## Load and Run Fine-Tuned Model

`RFDETR.from_checkpoint("<CHECKPOINT_PATH>")` picks the model class and restores the trained resolution and architecture when the checkpoint contains `model_config`. Best-total checkpoints written by RF-DETR 1.11.0 and earlier lack that metadata; load the matching unstripped `checkpoint_best_<source>.pth` when it is still beside the file, or pass the training architecture explicitly. `last.ckpt` and `checkpoint_<epoch>.ckpt` files written by RF-DETR 1.11.0 and earlier hold the weights but not the model class or its settings, so `from_checkpoint` rejects them. Load `checkpoint_best_ema.pth` from the same run instead, or pass the file as `pretrain_weights=` to the variant constructor it was trained with, together with the architecture settings listed in `training_config.json`; that route keeps the default class names, because those files do not store the dataset's. The variant constructors below read only the weights, so pass `resolution=` explicitly if you trained at a non-default resolution.

=== "Object Detection"

    ```python
    from rfdetr import RFDETRMedium

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

    detections = model.predict("<IMAGE_PATH>")
    ```

=== "Image Segmentation"

    ```python
    from rfdetr import RFDETRSegMedium

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

    detections = model.predict("<IMAGE_PATH>")
    ```

## Evaluate a Fine-Tuned Model

`model.evaluate()` runs a single evaluation pass over a dataset split and returns (and prints) the COCO metrics — mAP, mAR, and the macro-F1 sweep. It works both right after `model.train()` and on a model loaded from a checkpoint; the weights already in memory are evaluated, so no checkpoint file is re-loaded.

```python
from rfdetr import RFDETRMedium

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

metrics = model.evaluate(dataset_dir="<DATASET_PATH>", split="test")
print(metrics["test/mAP_50_95"])
```

- `split="test"` evaluates the `test/` folder on Roboflow-exported (`dataset_file="roboflow"`, the default) and YOLO (`dataset_file="yolo"`) datasets. YOLO-format datasets — plain YOLO plus Roboflow exports detected as YOLO — are not required to declare a `test` split; when it can't be resolved, `split="test"` falls back to the `valid/` folder instead and logs a warning — the returned metric keys are still prefixed `test/*`, so double-check the logs before treating `test/mAP_50_95` as held-out-test performance. A Roboflow export detected as COCO format has to ship `test/_annotations.coco.json`; without it `split="test"` raises `FileNotFoundError` rather than falling back. For COCO and Objects365 there is no dedicated test split at all, so `split="test"` silently evaluates the `valid/` folder instead. `split="val"` always evaluates `valid/` directly.
- The detection head is never adapted to the dataset — the model is evaluated exactly as configured. If the dataset's class count differs from the model's `num_classes`, a warning is emitted and evaluation proceeds unchanged.
- Evaluation writes no checkpoints or logs to `output_dir`.
- `evaluate()` accepts the same keyword arguments as `train()` for convenience, but training-only fields (`epochs`, `lr`, `ema`, `early_stopping`, logger flags, etc.) have no effect — evaluation runs through an eval-only trainer that never builds those callbacks.

## Next Steps

After training your model, you can:

- [Export your model to ONNX](../../exports/index.md) for deployment with various inference frameworks
- [Deploy to Roboflow](../deploy.md) for cloud-based inference and workflow integration