---
description: Customize RF-DETR training with PyTorch Lightning primitives. Direct access to RFDETRModelModule, RFDETRDataModule, and build_trainer.
---

# Custom Training API

The high-level `RFDETR.train()` method is the quickest path to fine-tuning, but the underlying training primitives are fully public and are the **recommended path for any customisation**: custom callbacks, alternative loggers, mixed-precision overrides, multi-GPU strategies, or integration with external training frameworks.

!!! tip "Quickstart vs. customisation"

    If you want to start training with minimal code, use `model.train()` — it sets up and runs the full PTL stack automatically. Come here when you need to take direct control over any part of that stack.

## How `RFDETR.train()` relates to PTL

When you call `model.train(...)`, three things happen internally:

```python
from rfdetr.training import RFDETRModelModule, RFDETRDataModule, build_trainer

module = RFDETRModelModule(model_config, train_config)
datamodule = RFDETRDataModule(model_config, train_config)
trainer = build_trainer(train_config, model_config)
trainer.fit(module, datamodule, ckpt_path=train_config.resume or None)
```

Each of these objects is a standard PTL class. You can construct them directly, modify them, swap out callbacks, or replace the trainer entirely.

---

## RFDETRModelModule

`RFDETRModelModule` is a `pytorch_lightning.LightningModule`. It owns the model weights, the criterion, the postprocessor, and the optimizer/scheduler configuration.

```python
from rfdetr.config import (
    RFDETRMediumConfig,
    TrainConfig,
)  # config classes live in rfdetr.config, not the top-level rfdetr namespace
from rfdetr.training import RFDETRModelModule

model_config = RFDETRMediumConfig(num_classes=10)
train_config = TrainConfig(
    dataset_dir="path/to/dataset",
    epochs=50,
    batch_size=4,
    grad_accum_steps=4,
    lr=1e-4,
    output_dir="output",
)

module = RFDETRModelModule(model_config, train_config)
```

### Lifecycle hooks

| Hook                       | Behaviour                                                                                                                                                                                                                       |
| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `on_fit_start`             | Seeds RNGs when `train_config.seed` is set.                                                                                                                                                                                     |
| `on_train_batch_start`     | Applies multi-scale random resize when `train_config.multi_scale=True`.                                                                                                                                                         |
| `transfer_batch_to_device` | Moves `NestedTensor` batches to the target device.                                                                                                                                                                              |
| `training_step`            | Computes loss and logs `train/loss` plus per-term losses. Keypoint models use manual optimization with box-normalized accumulation across microbatches; detection and segmentation use Lightning's automatic optimization path. |
| `validation_step`          | Runs forward pass and postprocessing; returns `{results, targets}` for `COCOEvalCallback`.                                                                                                                                      |
| `test_step`                | Same as `validation_step`, logs under `test/`.                                                                                                                                                                                  |
| `predict_step`             | Runs inference-only forward pass and returns postprocessed detections.                                                                                                                                                          |
| `configure_optimizers`     | Builds the configured optimizer, preserves layer-wise LR decay, applies optional parameter-group kwargs, and attaches the configured LR scheduler (managed preset, import path, or callable).                                   |
| `on_load_checkpoint`       | Auto-converts legacy `.pth` checkpoints to PTL format.                                                                                                                                                                          |

### Accessing the underlying model

The raw `nn.Module` is `module.model`. After training completes, `RFDETR.train()` syncs it back onto `self.model.model` so `predict()` and `export()` continue to work.

### Custom optimizer

`TrainConfig.optimizer` accepts either a **name / import path** (a string) or a **callable** (a class or `functools.partial`), and selects the optimizer built in `configure_optimizers`.

- **Managed short name** — a bare `torch.optim` optimizer name such as the default `"adamw"`, `"sgd"`, or `"adam"` (case-insensitive). **Only native `torch.optim` optimizers may be chosen by short name.** RF-DETR manages construction here: it injects `lr`, a signature-aware `weight_decay`, and your `optimizer_kwargs`.
- **Explicit import path** — a dotted string such as `"torch.optim.AdamW"` or `"pytorch_optimizer.Lion"`. This is how you select any non-`torch.optim` optimizer, e.g. one from [`pytorch-optimizer`](https://kozistr.tech/pytorch_optimizer/) (a third-party library you install yourself — RF-DETR does not depend on it). The class is constructed from the RF-DETR parameter groups (which already carry per-group learning rates) plus your `optimizer_kwargs` only; RF-DETR injects no `lr` or `weight_decay`, so pass any you need in `optimizer_kwargs`.
- **Callable** — a class or `functools.partial` is called as `optimizer(param_groups)` and nothing else, so bake all hyperparameters into the callable. `optimizer_kwargs` is ignored (with a warning) in this mode. A reconstructable callable (an importable top-level class / `functools.partial` with JSON-serializable keyword arguments) is desugared to its dotted import path plus keyword arguments so the config round-trips through `training_config.json`; a lambda, locally-defined class, or non-serializable argument still trains but cannot be restored from a saved config (RF-DETR warns with a specific hint).

```python
import functools

# managed native torch.optim short name — RF-DETR injects lr + weight_decay
train_config = TrainConfig(..., optimizer="adamw")

# explicit import path (incl. pytorch-optimizer) — supply hyperparameters via optimizer_kwargs
train_config = TrainConfig(..., optimizer="pytorch_optimizer.Lion", optimizer_kwargs={"weight_decouple": True})

# callable / functools.partial — bake arguments in; optimizer_kwargs is ignored
train_config = TrainConfig(..., optimizer=functools.partial(torch.optim.AdamW, weight_decay=1e-4))
```

See [Training parameters — optimizer](training-parameters.md) for the full parameter reference and examples.

`fused_optimizer` only affects the built-in `optimizer="adamw"` path; it is ignored for custom optimizers. On a BF16/CUDA run where fused AdamW would otherwise apply, RF-DETR logs a warning if `fused_optimizer=True` is combined with a non-default optimizer.

!!! warning "Wrapper optimizers (SAM, Lookahead, …) need extra care"

    Non-`torch.optim` optimizers are selected by import path (e.g. `"pytorch_optimizer.Lion"`) or a callable. Most `pytorch-optimizer` optimizers are drop-in replacements and work directly — including `Lion`, `AdaBelief`, `SophiaH`, `Adan`, `Ranger`, and `Ranger21`, which take `params` first like any standard optimizer. Two groups need extra care:

    - **Base-optimizer wrappers — `SAM` / `BSAM` / `GSAM` / `WSAM`.** They require a `base_optimizer` as a second positional argument, so `optimizer="pytorch_optimizer.SAM"` raises a `TypeError` at training start (RF-DETR calls the class with the parameter groups only). They also change `optimizer.step()` semantics (SAM needs two forward/backward passes per step) and are incompatible with PyTorch Lightning's default `automatic_optimization=True`. Supply the `base_optimizer` via a `functools.partial`, but such an optimizer will not round-trip through a saved config.
    - **`Lookahead` / `PCGrad` / `GradientCentralization`.** These wrap or post-process another optimizer and likewise cannot be built from parameter groups alone; construction with the parameter groups only raises a `TypeError`.

    If you need SAM or Lookahead, override `configure_optimizers` and set `self.automatic_optimization = False` in `RFDETRModelModule.__init__`. For standard use, pick a `torch.optim` optimizer by name or a drop-in `pytorch-optimizer` optimizer by import path (e.g. `"pytorch_optimizer.Lion"`).

### Custom LR scheduler

`TrainConfig.lr_scheduler` accepts either a **preset name / import path** (a string) or a **callable** (a class or `functools.partial`), mirroring `optimizer`, and selects the scheduler built in `configure_optimizers`.

- **Managed preset** — the built-in `"step"` (10x drop after `lr_drop` epochs) or `"cosine"` (annealing to `min_factor`). These own linear warmup and full-run step sizing; pass `lr_drop` / `min_factor` via `lr_scheduler_kwargs`.
- **Explicit import path** — a dotted string such as `"torch.optim.lr_scheduler.OneCycleLR"`. The class is constructed as `scheduler(optimizer, **lr_scheduler_kwargs)`; RF-DETR injects no `total_steps` / `T_max`, so pass any the scheduler needs in `lr_scheduler_kwargs`.
- **Callable** — a class or `functools.partial` is called as `lr_scheduler(optimizer)` and nothing else, so bake all hyperparameters into the callable. `lr_scheduler_kwargs` is ignored (with a warning) in this mode. Reconstructable callables desugar to a dotted path plus kwargs so the config round-trips through `training_config.json`; a lambda or locally-defined class still trains but cannot be restored from a saved config.

```python
import functools, torch

# managed preset — warmup + full-run sizing handled by RF-DETR
train_config = TrainConfig(..., lr_scheduler="cosine", lr_scheduler_kwargs={"min_factor": 0.1})

# explicit import path — supply scheduler arguments via lr_scheduler_kwargs
train_config = TrainConfig(
    ..., lr_scheduler="torch.optim.lr_scheduler.StepLR", lr_scheduler_kwargs={"step_size": 30, "gamma": 0.1}
)

# callable / functools.partial — bake arguments in; lr_scheduler_kwargs is ignored
train_config = TrainConfig(..., lr_scheduler=functools.partial(torch.optim.lr_scheduler.StepLR, step_size=30))
```

With `warmup_epochs > 0`, an explicit scheduler is automatically prepended with a linear warmup ramp via `SequentialLR` (managed presets bake warmup into their own schedule). `ReduceLROnPlateau` is special: it cannot be warmup-wrapped, always steps once per epoch, and reads the metric named by `lr_scheduler_monitor` (default `"val/loss"`). Use `lr_scheduler_interval="epoch"` to step other explicit schedulers per epoch instead of per optimizer step.

See [Training parameters — scheduler](training-parameters.md#scheduler-and-regularization) for the full parameter reference.

---

## RFDETRDataModule

`RFDETRDataModule` is a `pytorch_lightning.LightningDataModule`. It builds train/val/test datasets and wraps them in `DataLoader` objects.

```python
from rfdetr.training import RFDETRDataModule

datamodule = RFDETRDataModule(model_config, train_config)
```

### Stages

| Stage        | Datasets built                             |
| ------------ | ------------------------------------------ |
| `"fit"`      | `train` + `val`                            |
| `"validate"` | `val` only                                 |
| `"test"`     | `test` (or `val` for COCO-format datasets) |

The `setup(stage)` method is lazy — each split is built at most once, even if called multiple times.

### class_names property

```python
datamodule.setup("fit")
print(datamodule.class_names)  # e.g. ["cat", "dog", "person"]
```

Returns sorted category names from the COCO annotation file of the first available split, or `None` if the dataset has not been set up yet.

---

## build_trainer

`build_trainer` assembles a `pytorch_lightning.Trainer` with the full RF-DETR callback and logger stack. All `TrainConfig` fields are wired automatically.

```python
from rfdetr.training import build_trainer

trainer = build_trainer(train_config, model_config)
```

### What build_trainer configures

| Concern               | Source                                                                                                                                                                                                                                                         |
| --------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| Max epochs            | `train_config.epochs`                                                                                                                                                                                                                                          |
| Gradient accumulation | Detection/segmentation: `train_config.grad_accum_steps` forwarded to Trainer. Keypoint models: owned by `RFDETRModelModule` manual optimization (Trainer always sees `1`).                                                                                     |
| Gradient clipping     | Detection/segmentation: `train_config.clip_max_norm` forwarded to Trainer. Keypoint models: owned by `RFDETRModelModule` manual optimization (Trainer always sees `None`).                                                                                     |
| Mixed precision       | `model_config.amp` enables AMP; dtype resolved from `train_config.amp_dtype` (`"auto"` selects `bf16-mixed` on Ampere+, `"bf16"` / `"fp16"` force a specific dtype)                                                                                            |
| Accelerator           | `train_config.accelerator` (default `"auto"`)                                                                                                                                                                                                                  |
| Strategy              | Set via `train_config.strategy` (default `"auto"`) or pass `strategy=` as a `**trainer_kwarg` to `build_trainer`. Common values: `"auto"`, `"ddp"`, `"ddp_spawn"`. `TrainConfig` also exposes `devices` and `num_nodes` for multi-GPU and multi-node training. |
| Sync batch norm       | `train_config.sync_bn`                                                                                                                                                                                                                                         |
| Progress bar          | `train_config.progress_bar`                                                                                                                                                                                                                                    |
| Loggers               | CSVLogger always; TensorBoard, WandB, MLflow when their `train_config` flags are `True`                                                                                                                                                                        |
| Callbacks             | `RFDETREMACallback`, `DropPathCallback`, `COCOEvalCallback`, `BestModelCallback`, `RFDETREarlyStopping` (conditional)                                                                                                                                          |

### Overriding PTL Trainer kwargs

Pass keyword arguments accepted by `pytorch_lightning.Trainer` via `**trainer_kwargs`. Most keys override the built configuration.

**Detection and segmentation models** forward `accumulate_grad_batches` and `gradient_clip_val` to the Trainer normally — you can override them via `trainer_kwargs` or configure them on `TrainConfig` (`grad_accum_steps`, `clip_max_norm`).

**Keypoint models** use manual optimization, so `RFDETRModelModule` owns accumulation and clipping internally. `build_trainer()` forces `accumulate_grad_batches=1` and `gradient_clip_val=None` regardless of what is passed, and emits a `UserWarning` if those keys appear in `trainer_kwargs` so the override is visible:

```python
trainer = build_trainer(
    train_config,
    model_config,
    fast_dev_run=2,  # run 2 batches per epoch for a smoke test
    log_every_n_steps=10,
)
```

---

## Running the training loop

### Full training run

```python
from rfdetr.config import (
    RFDETRMediumConfig,
    TrainConfig,
)  # config classes live in rfdetr.config, not the top-level rfdetr namespace
from rfdetr.training import RFDETRModelModule, RFDETRDataModule, build_trainer

model_config = RFDETRMediumConfig(num_classes=10)
train_config = TrainConfig(
    dataset_dir="path/to/dataset",
    epochs=100,
    batch_size=4,
    grad_accum_steps=4,
    lr=1e-4,
    output_dir="output",
)

module = RFDETRModelModule(model_config, train_config)
datamodule = RFDETRDataModule(model_config, train_config)
trainer = build_trainer(train_config, model_config)

trainer.fit(module, datamodule)
```

### Resume from checkpoint

Pass the checkpoint path to `trainer.fit` via `ckpt_path`. The path can be a PTL `.ckpt` file or a legacy RF-DETR `.pth` file — `RFDETRModelModule.on_load_checkpoint` converts either format automatically.

```python
trainer.fit(module, datamodule, ckpt_path="output/last.ckpt")
# or a legacy checkpoint:
trainer.fit(module, datamodule, ckpt_path="output/checkpoint.pth")
```

> **Note:** When `checkpoint_interval=1`, no `last.ckpt` is written. Use `checkpoint_{epoch}.ckpt` (e.g. `output/checkpoint_epoch=4.ckpt`) to resume instead.

If you need to persist a converted checkpoint on disk (for example to inspect it, share it, or use it outside of PTL), convert it explicitly before passing it to `trainer.fit`:

```python
from rfdetr.training import convert_legacy_checkpoint

convert_legacy_checkpoint("old_checkpoint.pth", "new_checkpoint.ckpt")
trainer.fit(module, datamodule, ckpt_path="new_checkpoint.ckpt")
```

`convert_legacy_checkpoint` reads a pre-PTL `.pth` file produced by the legacy `engine.py` training loop and writes a PTL-compatible `.ckpt` file. Use it when migrating saved checkpoints to the PTL format rather than relying on on-the-fly conversion at load time.

### Validation only

```python
trainer.validate(module, datamodule)
```

Runs one full validation pass and logs `val/mAP_50_95`, `val/mAP_50`, `val/F1`, and per-class AP metrics to all active loggers.

### Inference with the data pipeline

```python
predictions = trainer.predict(module, dataloaders=datamodule.val_dataloader())
```

Calls `module.predict_step` on every batch and returns a list of postprocessed detection results. Pass any `DataLoader` instance — `datamodule.val_dataloader()`, `datamodule.test_dataloader()`, or a custom loader — as the `dataloaders` argument. This is useful for offline evaluation or generating submission files.

!!! note "predict_dataloader not implemented"

    `RFDETRDataModule` does not define a `predict_dataloader()` method, so `trainer.predict(module, datamodule)` will raise an error. Always pass a dataloader explicitly via the `dataloaders=` argument.

---

## Multi-GPU training

`build_trainer` configures PyTorch Lightning's `Trainer` directly, so all PTL strategies work out of the box.

### Data Parallel (DDP) — recommended

Set `train_config.accelerator = "auto"` and pass `strategy="ddp"` to `build_trainer`, then launch with `torchrun`:

!!! note "`devices` must be overridden for multi-GPU runs"

    `build_trainer` defaults to `devices=1`. To use all available GPUs, pass `devices="auto"` (or an explicit count) as a `**trainer_kwarg`:

    ```python
    trainer = build_trainer(train_config, model_config, strategy="ddp", devices="auto")
    ```

    Without this override, `torchrun` will spawn multiple processes but each process will only see one device, defeating the purpose of the multi-GPU launch.

```bash
torchrun --nproc_per_node=4 train.py
```

where `train.py` contains:

```python
from rfdetr.config import (
    RFDETRMediumConfig,
    TrainConfig,
)  # config classes live in rfdetr.config, not the top-level rfdetr namespace
from rfdetr.training import RFDETRModelModule, RFDETRDataModule, build_trainer

model_config = RFDETRMediumConfig(num_classes=10)
train_config = TrainConfig(
    dataset_dir="path/to/dataset",
    epochs=100,
    batch_size=4,  # per-GPU batch size
    grad_accum_steps=1,  # reduce when using more GPUs
    output_dir="output",
    sync_bn=True,  # sync batch norms across GPUs
)

module = RFDETRModelModule(model_config, train_config)
datamodule = RFDETRDataModule(model_config, train_config)
trainer = build_trainer(train_config, model_config, strategy="ddp", devices="auto")

trainer.fit(module, datamodule)
```

!!! warning "EMA is not compatible with FSDP or DeepSpeed"

    `build_trainer` automatically disables `RFDETREMACallback` when `strategy` contains `"fsdp"` or `"deepspeed"`, and emits a `UserWarning`. Use `strategy="ddp"` or `strategy="auto"` to keep EMA active.

### Effective batch size

```
effective_batch_size = batch_size × grad_accum_steps × num_gpus
```

Maintain an effective batch size of 16 regardless of GPU count:

| GPUs | `batch_size` | `grad_accum_steps` | Effective |
| ---- | ------------ | ------------------ | --------- |
| 1    | 4            | 4                  | 16        |
| 2    | 4            | 2                  | 16        |
| 4    | 4            | 1                  | 16        |
| 8    | 2            | 1                  | 16        |

---

## Custom callbacks

`build_trainer` builds the default callback stack. To add your own callbacks alongside the built-in ones, pass them via `trainer_kwargs`:

```python
from pytorch_lightning.callbacks import LearningRateMonitor, ModelSummary
from rfdetr.training import build_trainer

extra_callbacks = [
    LearningRateMonitor(logging_interval="step"),
    ModelSummary(max_depth=3),
]

trainer = build_trainer(
    train_config,
    model_config,
    callbacks=extra_callbacks,  # replaces the default callback list entirely
)
```

!!! warning "Replacing vs. extending callbacks"

    Passing `callbacks=` to `build_trainer` via `trainer_kwargs` **replaces** the entire default callback list built inside `build_trainer` (EMA, COCO eval, best-model checkpointing, etc.). To extend rather than replace, build the extra callbacks separately and merge them after calling `build_trainer`:

    ```python
    trainer = build_trainer(train_config, model_config)
    trainer.callbacks.extend(
        [
            LearningRateMonitor(logging_interval="step"),
        ]
    )
    trainer.fit(module, datamodule)
    ```

### Built-in callbacks

| Class                 | Purpose                                                                                     | Enabled when                                            |
| --------------------- | ------------------------------------------------------------------------------------------- | ------------------------------------------------------- |
| `RFDETREMACallback`   | Maintains an EMA copy of model weights                                                      | `train_config.use_ema=True` and strategy is not sharded |
| `DropPathCallback`    | Anneals drop-path rate over training                                                        | `train_config.drop_path > 0`                            |
| `COCOEvalCallback`    | Computes task validation metrics after each validation epoch                                | Always                                                  |
| `BestModelCallback`   | Saves `checkpoint_best_regular.pth`, `checkpoint_best_ema.pth`, `checkpoint_best_total.pth` | Always                                                  |
| `RFDETREarlyStopping` | Stops training when the validation task metric stops improving                              | `train_config.early_stopping=True`                      |

---

## Custom loggers

`build_trainer` adds loggers based on `TrainConfig` flags. To attach a logger not supported by `TrainConfig` (for example a custom Neptune or Comet logger), build it yourself and pass it alongside the defaults:

```python
from pytorch_lightning.loggers import NeptuneLogger  # hypothetical
from rfdetr.training import build_trainer

trainer = build_trainer(train_config, model_config)
trainer.loggers.append(NeptuneLogger(project="my-workspace/rf-detr"))
trainer.fit(module, datamodule)
```

All logged keys (`train/loss`, `val/mAP_50_95`, `val/keypoint_map_50_95`, `val/F1`, `val/ema_mAP_50_95`, etc.) are written to every active logger in the list.

---

## Logged metrics reference

| Key                      | When logged            | Description                                               |
| ------------------------ | ---------------------- | --------------------------------------------------------- |
| `train/loss`             | Every step / epoch     | Total weighted training loss                              |
| `train/<term>`           | Every step / epoch     | Individual loss terms (e.g. `train/loss_bbox`)            |
| `val/loss`               | Each epoch             | Validation loss (if `train_config.compute_val_loss=True`) |
| `val/mAP_50_95`          | Each eval epoch        | COCO box mAP@[.50:.05:.95]                                |
| `val/mAP_50`             | Each eval epoch        | COCO box mAP@.50                                          |
| `val/mAP_75`             | Each eval epoch        | COCO box mAP@.75                                          |
| `val/mAR`                | Each eval epoch        | COCO mean average recall                                  |
| `val/ema_mAP_50_95`      | Each eval epoch        | EMA-model mAP@[.50:.05:.95] (if EMA active)               |
| `val/F1`                 | Each eval epoch        | Macro F1 at best confidence threshold                     |
| `val/precision`          | Each eval epoch        | Precision at best F1 threshold                            |
| `val/recall`             | Each eval epoch        | Recall at best F1 threshold                               |
| `val/AP/<class>`         | Each eval epoch        | Per-class AP (if `log_per_class_metrics=True`)            |
| `val/segm_mAP_50_95`     | Each eval epoch        | Segmentation mAP (segmentation models only)               |
| `val/segm_mAP_50`        | Each eval epoch        | Segmentation mAP@.50 (segmentation models only)           |
| `val/keypoint_map_50_95` | Each eval epoch        | COCO keypoint AP@[.50:.05:.95] (keypoint preview only)    |
| `val/keypoint_map_50`    | Each eval epoch        | COCO keypoint AP@.50 (keypoint preview only)              |
| `test/*`                 | After `trainer.test()` | Mirror of `val/*` keys                                    |

---

## See also

- [RFDETR.train() — high-level API](index.md#quick-start) — the one-liner training path
- [Training parameters](training-parameters.md) — all `TrainConfig` fields
- [Training loggers](loggers.md) — TensorBoard, WandB, MLflow setup
- [Advanced training](advanced.md) — checkpointing, early stopping, memory optimisation
- [PTL primitives API reference](../../reference/training.md) — full docstring reference
