Training API Reference

This page documents the training primitives that power RF-DETR. For a narrative guide with runnable examples, see Custom Training API.

RFDETRModelModule

Bases: LightningModule

LightningModule wrapping the RF-DETR model and training loop.

Parameters:

Name	Type	Description	Default
model_config	ModelConfig	Architecture configuration.	required
train_config	TrainConfig	Training hyperparameter configuration.	required

Functions

__init__(model_config, train_config)

on_fit_start()

Seed RNGs at fit start when TrainConfig.seed is set.

This avoids hidden global side-effects in build_trainer while still preserving deterministic training behaviour for actual fit runs.

on_train_batch_start(batch, batch_idx)

Apply optional multi-scale resize to the incoming batch.

Modifications to batch (in-place on NestedTensor) are visible in training_step because they share the same object.

Parameters:

Name	Type	Description	Default
batch	tuple[Any, Any]	Tuple of (NestedTensor samples, list of target dicts).	required
batch_idx	int	Index of the current batch within the epoch.	required

training_step(batch, batch_idx)

Compute loss for one training step and log metrics.

PTL handles AMP (precision) without a manual GradScaler. Keypoint models perform manual optimization so box-count loss normalization is based on the full accumulated effective batch rather than each microbatch independently; detection and segmentation models keep Lightning's automatic optimization path.

Parameters:

Name	Type	Description	Default
batch	tuple[Any, Any]	Tuple of (NestedTensor samples, list of target dicts).	required
batch_idx	int	Batch index within the epoch.	required

Returns:

Type	Description
Tensor | dict[str, Any]	Scalar loss tensor by default. When compute_train_metrics=True,
Tensor | dict[str, Any]	returns a Lightning-compatible dict containing loss plus
Tensor | dict[str, Any]	detached postprocessed predictions for train mAP logging.

validation_step(batch, batch_idx)

Run forward pass and postprocess for one validation step.

Returns raw results and targets so COCOEvalCallback can accumulate them across the epoch via on_validation_batch_end.

Parameters:

Name	Type	Description	Default
batch	tuple[Any, Any]	Tuple of (NestedTensor samples, list of target dicts).	required
batch_idx	int	Batch index within the validation epoch.	required

Returns:

Type	Description
dict[str, Any]	Dict with results (postprocessed predictions) and targets.

test_step(batch, batch_idx)

Run forward pass and postprocess for one test step.

Mirrors :meth:validation_step so COCOEvalCallback can accumulate results via on_test_batch_end when trainer.test() is called (e.g. from :class:~rfdetr.training.callbacks.BestModelCallback at end of training).

Parameters:

Name	Type	Description	Default
batch	tuple[Any, Any]	Tuple of (NestedTensor samples, list of target dicts).	required
batch_idx	int	Batch index within the test epoch.	required

Returns:

Type	Description
dict[str, Any]	Dict with results (postprocessed predictions) and targets.

predict_step(batch, batch_idx, dataloader_idx=0)

Run inference on a preprocessed batch and return postprocessed results.

Parameters:

Name	Type	Description	Default
batch	tuple[Any, Any]	Tuple of (NestedTensor samples, list of target dicts).	required
batch_idx	int	Batch index.	required
dataloader_idx	int	Index of the predict dataloader.	0

Returns:

Type	Description
Any	Postprocessed detection results from PostProcess.

configure_optimizers()

Build the configured optimizer with layer-wise LR decay and scheduler.

Uses trainer.estimated_stepping_batches for total step count so cosine annealing covers the full training run regardless of dataset size or accumulation settings. optimizer="adamw" keeps RF-DETR's fused torch AdamW path; other names can be loaded from pytorch-optimizer.

Returns:

Type	Description
OptimizerLRSchedulerConfig	PTL optimizer config dict with optimizer and step-interval scheduler.

clip_gradients(optimizer, gradient_clip_val=None, gradient_clip_algorithm=None)

Override PTL gradient clipping to support fused AdamW.

PTL's AMP precision plugin refuses to clip gradients when the optimizer declares it handles unscaling internally (fused=True). When fused is active we are on BF16 (no GradScaler) so clip_grad_norm_ is correct. For the non-fused path (FP16 + GradScaler or FP32) we delegate to super() to preserve scaler-aware unscaling.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer | LightningOptimizer	The current optimizer.	required
gradient_clip_val	float | None	Maximum gradient norm.	None
gradient_clip_algorithm	str | None	Clipping algorithm; forwarded to super() for the non-fused path.	None

on_load_checkpoint(checkpoint)

Auto-detect legacy formats and reconcile PE shapes at checkpoint load time.

PTL calls this hook before applying checkpoint["state_dict"] to the module. Three normalisation steps are applied in order:

1. Raw legacy format — a *.pth file loaded directly by Trainer (e.g. via ckpt_path=). Recognised by the presence of "model" without "state_dict". The state dict is rewritten in-place with the "model." prefix so PTL can apply it normally.

2. Positional-embedding interpolation — when the checkpoint was saved at a different image resolution than the current model, the DINOv2 position_embeddings tensor shape will mismatch. :func:~rfdetr.models.weights.interpolate_position_embeddings is called to bicubic-resize the PE to model_config.positional_encoding_size before PTL applies the state dict. Regression fix for :issue:998.

3. Converted format — a file produced by :func:~rfdetr.training.checkpoint.convert_legacy_checkpoint that already has "state_dict" but also carries "legacy_ema_state_dict". The EMA weights are stashed on self._pending_legacy_ema_state for optional restoration by :class:~rfdetr.training.callbacks.ema.RFDETREMACallback.

Note

This hook only fires on Trainer(ckpt_path=...) resume paths. Fresh-train bootstrap from a pretrain_weights checkpoint runs through :func:~rfdetr.models.weights.load_pretrain_weights during __init__ instead — that helper performs its own PTL .ckpt normalisation (state_dict → model key, _orig_mod strip) and PE interpolation, so the two code paths intentionally do not share state.

Parameters:

Name	Type	Description	Default
checkpoint	dict[str, Any]	Checkpoint dict passed in by PTL (mutated in-place).	required

reinitialize_detection_head(num_classes)

Reinitialize the detection head for a new class count.

Parameters:

Name	Type	Description	Default
num_classes	int	New number of classes (excluding background).	required

---

RFDETRDataModule

Bases: LightningDataModule

LightningDataModule wrapping RF-DETR dataset construction and data loading.

Parameters:

Name	Type	Description	Default
model_config	ModelConfig	Architecture configuration (used for resolution, patch_size, etc.).	required
train_config	TrainConfig	Training hyperparameter configuration (used for dataset params).	required

Attributes

class_names property

Class names from the training or validation dataset annotation file.

Reads category names from the first available COCO-style dataset. Returns None if no dataset has been set up yet or the dataset does not expose COCO-style category information.

Returns:

Type	Description
list[str] | None	Sorted list of class name strings, or None.

Functions

__init__(model_config, train_config)

setup(stage)

Build datasets for the requested stage.

PTL calls this on every process before the corresponding dataloader method. Datasets are built lazily — a dataset is only constructed once even if setup is called multiple times.

Parameters:

Name	Type	Description	Default
stage	str	PTL stage identifier — one of "fit", "validate", "test", or "predict".	required

train_dataloader()

Return the training DataLoader.

Uses a replacement sampler when the dataset is too small to fill _MIN_TRAIN_BATCHES effective batches (matching legacy behaviour in main.py). Otherwise wraps the dataset with :class:GradAccumAlignedDataset to ensure its length is an exact multiple of effective_batch_size * world_size (workaround for https://github.com/Lightning-AI/pytorch-lightning/issues/19987) and then uses shuffle=True, drop_last=True so that PTL can auto-inject DistributedSampler in DDP mode.

Returns:

Type	Description
DataLoader[Any]	DataLoader for the training dataset.

val_dataloader()

Return the validation DataLoader.

Returns:

Type	Description
DataLoader[Any]	DataLoader for the validation dataset with sequential sampling.

test_dataloader()

Return the test DataLoader.

Returns:

Type	Description
DataLoader[Any]	DataLoader for the test dataset with sequential sampling.

---

build_trainer

Assemble a PTL Trainer with the full RF-DETR callback and logger stack.

Resolves training precision from model_config.amp and device capability, guards EMA against sharded strategies, wires conditional loggers, and applies promoted training knobs (sync_batchnorm, strategy).

Parameters:

Name	Type	Description	Default
train_config	TrainConfig	Training hyperparameter configuration.	required
model_config	ModelConfig	Architecture configuration. Used for precision resolution (model_config.amp) and to guard against unsupported distributed configurations for keypoint models.	required
accelerator	str | None	PTL accelerator string (e.g. "auto", "cpu", "gpu"). Defaults to None which reads from train_config.accelerator (itself defaulting to "auto"). Pass "cpu" to override auto-detection (e.g. when the caller explicitly requests CPU training via device="cpu").	None
include_training_callbacks	bool	When True (default) the full training stack is wired (EMA, drop-path, checkpointing, best-model selection, early stopping) along with the configured loggers. When False an evaluation-only trainer is built that keeps just the metric callback (and the progress bar): no checkpoints or logs are written. Used by :meth:rfdetr.detr.RFDETR.evaluate.	True
**trainer_kwargs	Any	Extra keyword arguments forwarded to pytorch_lightning.Trainer. Use this to pass PTL-native flags that are not exposed through TrainConfig, for example:: build_trainer(tc, mc, fast_dev_run=2) Most keys present in both trainer_kwargs and the built config dict are overridden by the value in trainer_kwargs. Detection and segmentation models forward accumulate_grad_batches from train_config.grad_accum_steps and gradient_clip_val from train_config.clip_max_norm to the Trainer normally. Keypoint models force accumulate_grad_batches=1 and gradient_clip_val=None because RFDETRModelModule owns both operations under manual optimization; passing those keys for a keypoint config raises a UserWarning to make the override explicit.	{}

Note

Two process-wide side effects: (1) unconditionally calls torch.set_float32_matmul_precision("high"), which persists after this function returns and overrides any caller-set precision (e.g. "highest") with no opt-out — mirrors the import-time guard in rfdetr.detr so the Lightning CLI path (rfdetr fit) gets the same TF32 behavior as the python API path. (2) sets check_val_every_n_epoch=tc.eval_interval, so eval_interval now gates the whole validation loop (forward pass, metric compute, EMA forward), not just result logging; a _ForceLastEpochValidationCallback still guarantees the final epoch always validates even when epochs is not a multiple of eval_interval.

When accelerator resolves to "xla" or "tpu", the resolved precision is set via an XLAPrecision plugin instead of a precision argument — PTL's XLAStrategy only accepts the plugin and raises TypeError for standard precision strings. The XLA plugin is appended to caller-supplied plugins.

Returns:

Type	Description
Trainer	A configured pytorch_lightning.Trainer instance.

---

Callbacks

RFDETREMACallback

Bases: Callback

Exponential Moving Average with optional tau-based warm-up.

Drop-in replacement for rfdetr.util.utils.ModelEma implemented as a plain Lightning callback around :class:torch.optim.swa_utils.AveragedModel. The _avg_fn reproduces the exact same formula as ModelEma (1-indexed updates counter, optional tau warm-up).

Parameters:

Name	Type	Description	Default
decay	float	Base EMA decay factor. Corresponds to TrainConfig.ema_decay.	0.993
tau	int	Warm-up time constant (in optimizer steps). When > 0 the effective decay ramps from 0 towards decay following decay * (1 - exp(-updates / tau)). Corresponds to TrainConfig.ema_tau.	100
use_buffers	bool	Whether buffers are averaged in addition to parameters.	True
update_interval_steps	int	Update EMA every N optimizer steps.	1

Attributes:

Name	Type	Description
suppress_test_swap		When True the test-epoch hooks skip the EMA weight swap. Set (and restored) by :class:~rfdetr.training.callbacks.best_model.BestModelCallback around its fit-end trainer.test() run, which has already loaded the best checkpoint weights into the module — swapping in the final EMA weights there would make the reported test/* metrics reflect the wrong model. Standalone trainer.test() runs keep the default False and evaluate EMA weights as before.

Functions

__init__(decay=0.993, tau=100, use_buffers=True, update_interval_steps=1)

BestModelCallback

Bases: ModelCheckpoint

Track best validation mAP and save best checkpoints during training.

Extends :class:pytorch_lightning.callbacks.ModelCheckpoint to save stripped {model, args, epoch} .pth files (instead of full .ckpt files) and to track a separate EMA checkpoint in parallel.

At the end of training the overall winner (regular vs EMA, strict > for EMA) is copied to checkpoint_best_total.pth and optimizer/scheduler state is stripped via :func:rfdetr.utilities.state_dict.strip_checkpoint. The stripped payload records best_total_source ("ema" or "regular") so the winning source is recoverable after reload.

When EMA tracking is enabled (monitor_ema set), EMA-named checkpoints are always left on disk for clarity: checkpoint_best_ema.pth (backfilled with the final EMA weights if the EMA metric never improved) and last_ema.pth (final EMA weights, mirroring last.pth for the live model).

Each track only updates when its own monitor key is actually logged that epoch, and the two tracks are checked independently. On non-eval epochs (eval_interval > 1 skips COCO eval entirely) both keys are absent and the callback is a full no-op. Under eval_ema_only (see TrainConfig.eval_ema_only), monitor_regular is never populated, so only the EMA track ever updates.

state_dict() and load_state_dict() are overridden to persist _best_ema in the Lightning callback state, ensuring that trainer.fit(ckpt_path=...) resumes EMA high-water-mark tracking from the correct value.

Parameters:

Name	Type	Description	Default
output_dir	str	Directory where checkpoint files are written.	required
monitor_regular	str	Metric key for the regular model mAP.	'val/mAP_50_95'
monitor_ema	str | None	Metric key for the EMA model mAP. None disables EMA tracking.	None
run_test	bool	If True, run trainer.test() on the best model at the end of training.	True
skip_best_epochs	int	Ignore the first N epochs (0..N-1) when tracking best regular and EMA checkpoints. Useful when fine-tuning from pretrain_weights: the pretrained model's epoch-0 mAP can artificially dominate best-checkpoint selection before training adapts to the new dataset.	0
smooth_alpha	float	Exponential-moving-average smoothing factor in [0.0, 1.0) applied to the regular monitor metric before checkpoint comparison. 0.0 (default) disables smoothing and keeps legacy behaviour: trainer.callback_metrics[monitor_regular] is consumed as-is by the parent :class:~pytorch_lightning.callbacks.ModelCheckpoint. smooth_alpha > 0 maintains an internal EMA state self._smoothed_regular = alpha * self._smoothed_regular + (1 - alpha) * raw and temporarily substitutes the smoothed value into trainer.callback_metrics for the duration of the parent's improvement check; the original raw value is always restored before returning so what gets logged to metrics.csv is unaffected. Useful for noisy metrics (e.g. keypoint mAP under NLL-Cholesky losses) where raw per-epoch swings can lock the best checkpoint to an early local peak. The EMA accumulator is updated on every validation epoch including epochs within the skip_best_epochs window so the smoothed value is pre-warmed by the first eligible comparison.	0.0

Examples:

Skip the first 3 epochs so pretrained weights do not dominate:

>>> import tempfile
>>> from rfdetr.training.callbacks.best_model import BestModelCallback
>>> with tempfile.TemporaryDirectory() as tmp:
...     cb = BestModelCallback(output_dir=tmp, skip_best_epochs=3)
...     cb._skip_best_epochs
3

Functions

__init__(output_dir, monitor_regular='val/mAP_50_95', monitor_ema=None, run_test=True, skip_best_epochs=0, smooth_alpha=0.0)

RFDETREarlyStopping

Bases: EarlyStopping

Early stopping callback monitoring validation mAP for RF-DETR.

Extends :class:pytorch_lightning.callbacks.EarlyStopping with dual-metric monitoring: by default it monitors max(regular_mAP, ema_mAP) (legacy behaviour); set use_ema=True to monitor the EMA metric exclusively.

The effective metric is injected into trainer.callback_metrics under a synthetic key before delegating to the parent's stopping logic, so all parent features are available for free: state_dict/load_state_dict for checkpoint resumption, NaN/inf guard via check_finite, and stopping_threshold/divergence_threshold.

Early stopping evaluates only on validation epochs where the monitored metrics are logged; non-eval epochs (eval_interval > 1) are skipped automatically.

Parameters:

Name	Type	Description	Default
patience	int	Number of epochs with no improvement before stopping.	10
min_delta	float	Minimum mAP improvement to reset the patience counter.	0.001
use_ema	bool	When True and both regular and EMA metrics are available, monitor only the EMA metric. When False, monitor max(regular, ema).	False
monitor_regular	str	Metric key for the regular model mAP.	'val/mAP_50_95'
monitor_ema	str	Metric key for the EMA model mAP.	'val/ema_mAP_50_95'
verbose	bool	If True, log early stopping status each epoch.	True
skip_best_epochs	int	Ignore the first N epochs (0..N-1) when evaluating patience and best-score baselines. Set this when fine-tuning from pretrain_weights to avoid premature stopping before the model adapts to the new dataset.	0

Examples:

Fine-tuning from pretrained weights — skip first 3 epochs:

>>> from rfdetr.training.callbacks.best_model import RFDETREarlyStopping
>>> cb = RFDETREarlyStopping(patience=10, skip_best_epochs=3)
>>> cb._skip_best_epochs
3

Functions

__init__(patience=10, min_delta=0.001, use_ema=False, monitor_regular='val/mAP_50_95', monitor_ema='val/ema_mAP_50_95', verbose=True, skip_best_epochs=0)

DropPathCallback

Bases: Callback

Applies per-step drop-path and dropout rate schedules to the model.

Computes the full schedule array in on_train_start using :func:rfdetr.util.drop_scheduler.drop_scheduler, then indexes into it on every training batch to update the model's stochastic-depth and dropout rates.

Parameters:

Name	Type	Description	Default
drop_path	float	Peak drop-path rate. 0.0 disables the schedule.	0.0
dropout	float	Peak dropout rate. 0.0 disables the schedule.	0.0
cutoff_epoch	int	Epoch boundary for early / late modes.	0
mode	Literal['standard', 'early', 'late']	Schedule mode forwarded to drop_scheduler.	'standard'
schedule	Literal['constant', 'linear']	Schedule shape forwarded to drop_scheduler.	'constant'
vit_encoder_num_layers	int	Passed to model.update_drop_path so the model can distribute rates across ViT encoder layers.	12

Functions

__init__(drop_path=0.0, dropout=0.0, cutoff_epoch=0, mode='standard', schedule='constant', vit_encoder_num_layers=12)

COCOEvalCallback

Bases: Callback

Validation callback that computes mAP (via torchmetrics) and macro-F1.

Accumulates predictions and targets across validation batches, then at epoch end computes:

• val/mAP_50_95, val/mAP_50, val/mAP_75, val/mAR using torchmetrics.detection.MeanAveragePrecision.
• Per-class val/AP/<name> when class names are available.
• val/F1, val/precision, val/recall from a confidence-threshold sweep over compact per-class matching data (DDP-safe).

For segmentation models (segmentation=True) additional metrics val/segm_mAP_50_95 and val/segm_mAP_50 are logged.

Parameters:

Name	Type	Description	Default
max_dets	int	Maximum detections per image passed to MeanAveragePrecision. Defaults to :data:~rfdetr.evaluation.keypoint_oks.DEFAULT_KEYPOINT_MAX_DETS.	DEFAULT_KEYPOINT_MAX_DETS
segmentation	bool	When True, evaluate both bbox and segm IoU using backend="faster_coco_eval". Defaults to False.	False
eval_interval	int	Run validation metrics every N epochs. Test metrics are always computed when trainer.test() is called.	1
log_per_class_metrics	bool	When False, skip per-class AP computation (MeanAveragePrecision(class_metrics=False)) as well as the per-class logging/table.	True
eval_ema_only	bool	When True, validation_step already forwarded through the EMA model directly (see TrainConfig.eval_ema_only), so the independent duplicate EMA forward pass this callback would otherwise run every validation batch is skipped.	False

Functions

__init__(max_dets=DEFAULT_KEYPOINT_MAX_DETS, segmentation=False, eval_interval=1, log_per_class_metrics=True, keypoint_oks_sigmas=None, in_notebook=None, eval_ema_only=False)

GPUMemoryTQDMProgressBar

Bases: _GpuMemoryMetricsMixin, TQDMProgressBar

TQDMProgressBar with peak CUDA memory usage in the postfix.

Functions

get_metrics(trainer, pl_module)

Return the progress bar metrics, with max_mem appended on CUDA.

An explicitly logged max_mem wins: if the base metrics already carry that key with a different value (a user's own self.log("max_mem", ..., prog_bar=True)), the logged value is kept and a rank-zero warning is emitted, mirroring how ProgressBar.get_metrics reports name collisions.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The Lightning Trainer instance.	required
pl_module	LightningModule	The LightningModule being trained.	required

Returns:

Type	Description
_Metrics	The metrics dict produced by the base progress bar, plus max_mem when
_Metrics	trainer.strategy.root_device is an active CUDA device.

GPUMemoryRichProgressBar

Bases: _GpuMemoryMetricsMixin, RichProgressBar

RichProgressBar with peak CUDA memory usage in the metrics column.

Functions

get_metrics(trainer, pl_module)

Return the progress bar metrics, with max_mem appended on CUDA.

An explicitly logged max_mem wins: if the base metrics already carry that key with a different value (a user's own self.log("max_mem", ..., prog_bar=True)), the logged value is kept and a rank-zero warning is emitted, mirroring how ProgressBar.get_metrics reports name collisions.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The Lightning Trainer instance.	required
pl_module	LightningModule	The LightningModule being trained.	required

Returns:

Type	Description
_Metrics	The metrics dict produced by the base progress bar, plus max_mem when
_Metrics	trainer.strategy.root_device is an active CUDA device.

---

RFDETRCli

CLI requires the train and cli extras

pip install "rfdetr[train,cli]"

The rfdetr console script and its --config / --print_config flags are provided by jsonargparse, which is included in the cli extra.

RFDETRCli is the command-line entry point for RF-DETR. It wraps RFDETRModelModule and RFDETRDataModule under a single rfdetr command and auto-generates four subcommands from the PyTorch Lightning CLI machinery:

rfdetr fit      --config configs/rfdetr_base.yaml
rfdetr validate --ckpt_path output/best.ckpt
rfdetr test     --ckpt_path output/best.ckpt
rfdetr predict  --ckpt_path output/best.ckpt

Both model_config and train_config are specified once; RFDETRCli automatically links them to the datamodule so you do not need to repeat the same arguments under --data.*.

Bases: LightningCLI

LightningCLI subclass for RF-DETR training and evaluation.

Wires RFDETRModelModule and RFDETRDataModule under a unified CLI, with argument linking that shares model_config and train_config between module and datamodule so the user only specifies them once.

Auto-generated subcommands: fit, validate, test, predict.