| # 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, divides by `accumulate_grad_batches`, and logs `train/loss` and per-term losses. | | |
| | `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 AdamW with layer-wise LR decay and a LambdaLR scheduler (cosine or step). | | |
| | `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. | |
| --- | |
| ## 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 | `train_config.grad_accum_steps` | | |
| | Gradient clipping | `train_config.clip_max_norm` (default `0.1`) | | |
| | Mixed precision | Resolved from `model_config.amp` and device capability (`bf16-mixed` on Ampere+, `16-mixed` otherwise) | | |
| | Accelerator | `train_config.accelerator` (default `"auto"`) | | |
| | Strategy | Pass `strategy=` as a `**trainer_kwarg` to `build_trainer`. `TrainConfig` has no `strategy` field — setting it on `TrainConfig` will raise a `ValueError`. | | |
| | 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 any keyword argument accepted by `pytorch_lightning.Trainer` via `**trainer_kwargs`. These override the built configuration: | |
| ```python | |
| trainer = build_trainer( | |
| train_config, | |
| model_config, | |
| fast_dev_run=2, # run 2 batches per epoch for a smoke test | |
| accumulate_grad_batches=8, # override TrainConfig.grad_accum_steps | |
| 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 mAP and F1 after each validation epoch | Always | | |
| | `BestModelCallback` | Saves `checkpoint_best_regular.pth`, `checkpoint_best_ema.pth`, `checkpoint_best_total.pth` | Always | | |
| | `RFDETREarlyStopping` | Stops training when validation mAP 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/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) | | |
| | `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 | |