wehub-resource-sync
16031aae96
CPU tests Workflow / Testing (ubuntu-latest, 3.12) (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.13) (push) Failing after 0s
Mypy Type Check / Type Check (push) Failing after 0s
Docs/Test WorkFlow / Test docs build (push) Failing after 1s
PR Conflict Labeler / labeling (push) Failing after 1s
Dependency resolution / Resolve [tflite] extra — Python 3.12 (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.10) (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.13) (push) Failing after 1s
CPU tests Workflow / build-pkg (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.10) (push) Failing after 0s
CPU tests Workflow / Testing (ubuntu-latest, 3.11) (push) Failing after 0s
Smoke Tests / try-all-models (macos-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (macos-latest, 3.13) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / testing-guardian (push) Has been cancelled
GPU tests Workflow / Testing (push) Has been cancelled
385 lines
21 KiB
Markdown
385 lines
21 KiB
Markdown
---
|
||
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 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 | 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
|