chore: import upstream snapshot with attribution

2026-07-13 12:26:24 +08:00
commit 16031aae96
343 changed files with 88674 additions and 0 deletions
@@ -0,0 +1,381 @@
+---
+description: Advanced RF-DETR training with resume, early stopping, multi-GPU DDP, gradient checkpointing, and memory optimization for large models.
+---
+
+# Advanced Training
+
+This page covers advanced training topics including resuming training, early stopping, multi-GPU training, and memory optimization techniques.
+
+!!! tip "PTL API for deeper customisation"
+
+    All examples on this page use the `RFDETR.train()` high-level API. For custom callbacks, non-default loggers, and fine-grained distributed training control, see the [Custom Training API](customization.md) guide.
+
+## Resume Training
+
+You can resume training from a previously saved checkpoint by passing the path to the `checkpoint.pth` file using the `resume` argument. This is useful when training is interrupted or you want to continue fine-tuning an already partially trained model.
+
+The training loop will automatically load:
+
+- Model weights
+- Optimizer state
+- Learning rate scheduler state
+- Training epoch number
+
+=== "Object Detection"
+
+    ```python
+    from rfdetr import RFDETRMedium
+
+    model = RFDETRMedium()
+
+    model.train(
+        dataset_dir="path/to/dataset",
+        epochs=100,
+        batch_size=4,
+        grad_accum_steps=4,
+        lr=1e-4,
+        output_dir="output",
+        resume="output/checkpoint.pth",
+    )
+    ```
+
+=== "Image Segmentation"
+
+    ```python
+    from rfdetr import RFDETRSegMedium
+
+    model = RFDETRSegMedium()
+
+    model.train(
+        dataset_dir="path/to/dataset",
+        epochs=100,
+        batch_size=4,
+        grad_accum_steps=4,
+        lr=1e-4,
+        output_dir="output",
+        resume="output/checkpoint.pth",
+    )
+    ```
+
+!!! tip "Resume vs Pretrain Weights"
+
+    - Use `resume="checkpoint.pth"` to continue training with optimizer state
+    - Use `pretrain_weights="checkpoint_best_total.pth"` when initializing a model to start fresh training from those weights
+
+---
+
+## Early Stopping
+
+Early stopping monitors the validation task metric and halts training if improvements remain below a threshold for a
+set number of epochs. Detection and segmentation models use box mAP; keypoint preview models use COCO keypoint AP.
+
+### Basic Usage
+
+=== "Object Detection"
+
+    ```python
+    from rfdetr import RFDETRMedium
+
+    model = RFDETRMedium()
+
+    model.train(
+        dataset_dir="path/to/dataset",
+        epochs=100,
+        batch_size=4,
+        grad_accum_steps=4,
+        lr=1e-4,
+        output_dir="output",
+        early_stopping=True,
+    )
+    ```
+
+=== "Image Segmentation"
+
+    ```python
+    from rfdetr import RFDETRSegMedium
+
+    model = RFDETRSegMedium()
+
+    model.train(
+        dataset_dir="path/to/dataset",
+        epochs=100,
+        batch_size=4,
+        grad_accum_steps=4,
+        lr=1e-4,
+        output_dir="output",
+        early_stopping=True,
+    )
+    ```
+
+### Configuration Options
+
+| Parameter                  | Default | Description                                          |
+| -------------------------- | ------- | ---------------------------------------------------- |
+| `early_stopping_patience`  | 10      | Number of epochs without improvement before stopping |
+| `early_stopping_min_delta` | 0.001   | Minimum metric change to count as improvement        |
+| `early_stopping_use_ema`   | False   | Use EMA model metrics for comparisons                |
+
+### Advanced Example
+
+```python
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=200,
+    early_stopping=True,
+    early_stopping_patience=15,  # Wait 15 epochs before stopping
+    early_stopping_min_delta=0.005,  # Require 0.5% validation metric improvement
+    early_stopping_use_ema=True,  # Track EMA model performance
+)
+```
+
+### How It Works
+
+1. After each epoch, the validation task metric is computed
+2. If the metric improves by at least `min_delta`, the patience counter resets
+3. If the metric doesn't improve, the patience counter increments
+4. When patience counter reaches `patience`, training stops
+5. The best checkpoint is already saved as `checkpoint_best_total.pth`
+
+```
+Epoch 10: mAP = 0.450 (best: 0.450) - counter: 0
+Epoch 11: mAP = 0.455 (best: 0.455) - counter: 0 (improved)
+Epoch 12: mAP = 0.454 (best: 0.455) - counter: 1 (no improvement)
+Epoch 13: mAP = 0.453 (best: 0.455) - counter: 2
+...
+Epoch 22: mAP = 0.452 (best: 0.455) - counter: 10 → STOP
+```
+
+---
+
+## Multi-GPU Training
+
+RF-DETR's training stack is built on PyTorch Lightning, so multi-GPU and multi-node training use the Lightning `Trainer` strategies directly. You can start multi-GPU runs through the high-level API or by using the Lightning primitives explicitly.
+
+### Using RFDETR.train() with multiple GPUs
+
+Create a training script and launch it with `torchrun`:
+
+```python
+# train.py
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium()
+
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,  # per-GPU batch size
+    grad_accum_steps=1,
+    lr=1e-4,
+    output_dir="output",
+    devices="auto",  # required — see note below
+)
+```
+
+```bash
+torchrun --nproc_per_node=4 train.py
+```
+
+!!! warning "Pass `devices=` explicitly"
+
+    `build_trainer()` defaults to `devices=1`. Without overriding this, training silently
+    runs on a single GPU even when `torchrun` launches multiple processes.
+
+    Pass `devices="auto"` to use all GPUs visible to the process, or pass an explicit
+    integer (e.g. `devices=4`). These values are forwarded to `build_trainer` via
+    `**trainer_kwargs`:
+
+    ```python
+    model.train(
+        dataset_dir="path/to/dataset",
+        epochs=100,
+        batch_size=4,
+        grad_accum_steps=1,
+        lr=1e-4,
+        output_dir="output",
+        devices="auto",  # or devices=4
+    )
+    ```
+
+### Batch Size with Multiple GPUs
+
+When using multiple GPUs, your effective batch size is multiplied by the number of GPUs:
+
+```
+effective_batch_size = batch_size × grad_accum_steps × num_gpus
+```
+
+**Example configurations for effective batch size of 16:**
+
+| GPUs | `batch_size` | `grad_accum_steps` | Effective |
+| ---- | ------------ | ------------------ | --------- |
+| 1    | 4            | 4                  | 16        |
+| 2    | 4            | 2                  | 16        |
+| 4    | 4            | 1                  | 16        |
+| 8    | 2            | 1                  | 16        |
+
+!!! warning "Adjust for GPU count"
+
+    When switching between single and multi-GPU training, remember to adjust `batch_size` and `grad_accum_steps` to maintain the same effective batch size.
+
+### Multi-Node Training
+
+For training across multiple machines, pass the standard `torchrun` flags:
+
+```bash
+torchrun \
+    --nproc_per_node=8 \
+    --nnodes=2 \
+    --node_rank=0 \
+    --master_addr="192.168.1.1" \
+    --master_port=1234 \
+    train.py
+```
+
+Run this command on each node, changing `--node_rank` accordingly.
+
+### Advanced multi-GPU options (PTL API)
+
+For fine-grained control over strategy, sync batch norm, precision, and other distributed settings, use the Lightning API directly.
+
+→ **[Multi-GPU with the PTL API](customization.md#multi-gpu-training)**
+
+---
+
+## Custom Augmentations
+
+RF-DETR supports advanced data augmentations using the [Albumentations](https://albumentations.ai/) library, providing access to over 70 different image transformations optimized for object detection.
+
+→ **[Complete Augmentation Guide](augmentations.md)** - Configuration examples, best practices, troubleshooting, and advanced topics.
+
+### Quick Start
+
+Pass an `aug_config` dictionary to `model.train()`. Each key is an Albumentations transform name; the value is a dict of keyword arguments for that transform:
+
+```python
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium()
+
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,
+    grad_accum_steps=4,
+    lr=1e-4,
+    output_dir="output",
+    aug_config={
+        "HorizontalFlip": {"p": 0.5},
+        "VerticalFlip": {"p": 0.5},
+        "Rotate": {"limit": 45, "p": 0.5},
+    },
+)
+```
+
+Use a built-in preset by importing it from `rfdetr.datasets.aug_configs`:
+
+```python
+from rfdetr.datasets.aug_configs import AUG_CONSERVATIVE, AUG_AGGRESSIVE, AUG_AERIAL, AUG_INDUSTRIAL
+
+model.train(dataset_dir="path/to/dataset", aug_config=AUG_AGGRESSIVE)
+```
+
+To disable all augmentations, pass an empty dict:
+
+```python
+model.train(dataset_dir="path/to/dataset", aug_config={})
+```
+
+---
+
+## Memory Optimization
+
+### Gradient Checkpointing
+
+For large models or high resolutions, enable gradient checkpointing to trade compute for memory.
+
+!!! warning "Constructor parameter — not a `train()` parameter"
+
+    `gradient_checkpointing` is a `ModelConfig` field and must be passed to the **model constructor**, not to `train()`. Passing it to `train()` will raise a `ValidationError` because `TrainConfig` has `extra="forbid"`.
+
+```python
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium(gradient_checkpointing=True)
+
+model.train(
+    dataset_dir="path/to/dataset",
+    batch_size=2,  # May be able to increase with checkpointing
+)
+```
+
+This re-computes activations during the backward pass instead of storing them, reducing memory usage by ~30-40% at the cost of ~20% slower training.
+
+### Memory-Efficient Configurations
+
+| Memory Level      | Configuration                                                                          |
+| ----------------- | -------------------------------------------------------------------------------------- |
+| Very Low (8GB)    | `batch_size=1`, `grad_accum_steps=16`, `gradient_checkpointing=True`, `resolution=576` |
+| Low (12GB)        | `batch_size=2`, `grad_accum_steps=8`, `gradient_checkpointing=True`                    |
+| Medium (16GB)     | `batch_size=4`, `grad_accum_steps=4`                                                   |
+| High (24GB)       | `batch_size=8`, `grad_accum_steps=2`                                                   |
+| Very High (40GB+) | `batch_size=16`, `grad_accum_steps=1`, `resolution=768`                                |
+
+---
+
+## Training Tips
+
+### Learning Rate Tuning
+
+- **Fine-tuning from COCO weights (default):** Use default learning rates (`lr=1e-4`, `lr_encoder=1.5e-4`)
+- **Small dataset (\<1000 images):** Consider lower `lr` (e.g., `5e-5`) to prevent overfitting
+- **Large dataset (>10000 images):** May benefit from higher `lr` (e.g., `2e-4`)
+
+### Epoch Count
+
+| Dataset Size      | Recommended Epochs |
+| ----------------- | ------------------ |
+| < 500 images      | 100-200            |
+| 500-2000 images   | 50-100             |
+| 2000-10000 images | 30-50              |
+| > 10000 images    | 20-30              |
+
+Use early stopping to automatically determine the optimal stopping point.
+
+### Data Augmentation
+
+RF-DETR applies built-in augmentations during training:
+
+- Random resizing
+- Random cropping
+- Color jittering
+- Horizontal flipping
+
+These are automatically configured and don't require manual setup.
+
+---
+
+## Troubleshooting
+
+### Out of Memory (OOM)
+
+If you encounter CUDA out of memory errors:
+
+1. Reduce `batch_size`
+2. Enable `gradient_checkpointing=True` (pass to the model constructor, not `train()`)
+3. Reduce `resolution`
+4. Increase `grad_accum_steps` to maintain effective batch size
+
+### Training Too Slow
+
+1. Increase `batch_size` (if memory allows)
+2. Use multiple GPUs with DDP
+3. Ensure you're using GPU (check `device="cuda"`)
+4. Consider using a smaller model (e.g., `RFDETRSmall` instead of `RFDETRLarge`)
+
+### Loss Not Decreasing
+
+1. Check that your dataset is correctly formatted
+2. Verify annotations are correct (bounding boxes in correct format)
+3. Try reducing the learning rate
+4. Check for class imbalance in your dataset
@@ -0,0 +1,156 @@
+---
+description: Configure RF-DETR data augmentations with Albumentations. Built-in presets for aerial, industrial, and small datasets plus custom transforms.
+---
+
+# Augmentations
+
+RF-DETR supports custom data augmentations via [Albumentations](https://albumentations.ai/), with automatic bounding box and mask handling for geometric transforms. Albumentations 1.4.24+ and 2.x are supported.
+
+## Quick Start
+
+Pass `aug_config` to your training call. Import one of the built-in presets:
+
+```python
+from rfdetr import RFDETRSmall
+from rfdetr.datasets.aug_configs import AUG_CONSERVATIVE, AUG_AGGRESSIVE, AUG_AERIAL, AUG_INDUSTRIAL
+
+model = RFDETRSmall()
+model.train(dataset_dir="path/to/dataset", epochs=100, aug_config=AUG_CONSERVATIVE)
+```
+
+Or pass a custom dict directly — keys are Albumentations transform names:
+
+```python
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    aug_config={
+        "HorizontalFlip": {"p": 0.5},
+        "Rotate": {"limit": 15, "p": 0.3},
+        "GaussianBlur": {"p": 0.2},
+    },
+)
+```
+
+To disable augmentations: `aug_config={}`. Omitting it uses the default (horizontal flip at 50%).
+
+## Built-in Presets
+
+| Preset             | Best for                          |
+| ------------------ | --------------------------------- |
+| `AUG_CONSERVATIVE` | Small datasets (under 500 images) |
+| `AUG_AGGRESSIVE`   | Large datasets (2000+ images)     |
+| `AUG_AERIAL`       | Satellite / overhead imagery      |
+| `AUG_INDUSTRIAL`   | Manufacturing / inspection data   |
+
+All presets are plain dicts — inspect or extend them before passing:
+
+```python
+from rfdetr.datasets.aug_configs import AUG_AGGRESSIVE
+
+my_config = {**AUG_AGGRESSIVE, "VerticalFlip": {"p": 0.1}}
+model.train(dataset_dir="...", aug_config=my_config)
+```
+
+### Recommendations by Dataset Size
+
+| Dataset Size     | Recommended preset                                              |
+| ---------------- | --------------------------------------------------------------- |
+| Under 500 images | `AUG_CONSERVATIVE` — flip + mild brightness/contrast            |
+| 500–2000 images  | Default or `AUG_CONSERVATIVE` with a few extra transforms added |
+| 2000+ images     | `AUG_AGGRESSIVE` — rotations, affine, color jitter              |
+
+## Nested Transforms
+
+RF-DETR supports `OneOf`, `SomeOf`, and `Sequential` container transforms from Albumentations. The most common pattern is `OneOf`, which randomly picks one transform from a group:
+
+```python
+aug_config = {
+    "HorizontalFlip": {"p": 0.5},
+    "OneOf": {
+        "transforms": [
+            {"Rotate": {"limit": 45, "p": 1.0}},
+            {"Affine": {"scale": (0.8, 1.2), "p": 1.0}},
+        ],
+    },
+    "GaussianBlur": {"p": 0.2},
+}
+```
+
+Each child's `p` controls its relative selection weight. The container itself always fires.
+
+If you need the same transform twice, or want explicit ordering, pass a list instead of a dict:
+
+```python
+aug_config = [
+    {"HorizontalFlip": {"p": 0.5}},
+    {"Rotate": {"limit": 45, "p": 0.3}},
+    {"Rotate": {"limit": 5, "p": 0.5}},  # second Rotate — only possible with list format
+]
+```
+
+Bounding boxes are updated automatically when a container holds any geometric transform — no extra configuration needed.
+
+## Geometric vs. Pixel-Level Transforms
+
+RF-DETR automatically handles bounding boxes for **geometric transforms** (flips, rotations, crops, affine, perspective). **Pixel-level transforms** (blur, noise, color) preserve coordinates unchanged. You don't need to handle this distinction — it's automatic based on the transform name.
+
+## Best Practices
+
+!!! tip "Start Conservative"
+
+    Begin with simple augmentations (horizontal flip, small brightness changes) and gradually add more as needed.
+
+!!! warning "Geometric Transforms"
+
+    Be careful with aggressive rotations and crops on datasets where object orientation matters (e.g., text detection, oriented objects).
+
+- **CPU-bound:** Augmentations run on CPU during data loading — more transforms means slower loading
+- **Use `num_workers`:** Parallelize augmentation across data loader workers
+- **Monitor training mAP vs validation mAP:** With strong augmentations it's normal for training mAP to be lower — validation uses original images while training uses augmented (harder) ones
+
+## Troubleshooting
+
+**Training is slow** — reduce the number of transforms or increase `num_workers`.
+
+**Boxes disappear after augmentation** — aggressive rotations or crops can push boxes outside the image boundary. Reduce rotation angles or avoid large crops.
+
+**Model not improving** — augmentations may be too aggressive. Start with `AUG_CONSERVATIVE` and add transforms gradually. Try removing geometric transforms first to isolate the cause.
+
+**Validation mAP is much higher than training mAP** — this is expected with strong augmentations and not a bug. See the monitoring tip above.
+
+**Upgrading albumentations to 2.x with existing `RandomSizedCrop` configs?** RF-DETR automatically adapts `height`/`width` kwargs to the `size=(height, width)` format required by albumentations 2.x. No config changes needed.
+
+## Advanced: Custom Transforms
+
+Any Albumentations transform works by name. If your custom transform is geometric, register it in `rfdetr/datasets/transforms.py` so boxes are updated automatically:
+
+```python
+GEOMETRIC_TRANSFORMS = {
+    ...,
+    "YourCustomTransform",
+}
+```
+
+Then use it like any other transform:
+
+```python
+model.train(
+    dataset_dir="...",
+    aug_config={
+        "HorizontalFlip": {"p": 0.5},
+        "YourCustomTransform": {"param": 1, "p": 0.3},
+    },
+)
+```
+
+## Reference
+
+- [Albumentations docs](https://albumentations.ai/docs/)
+- [All available transforms](https://albumentations.ai/docs/api_reference/augmentations/)
+
+## Next Steps
+
+- [Monitor training with TensorBoard](loggers.md#tensorboard)
+- [Use early stopping](advanced.md#early-stopping) to prevent overfitting
+- [Export your trained model](../export.md) for deployment
@@ -0,0 +1,384 @@
+---
+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
@@ -0,0 +1,493 @@
+---
+description: RF-DETR dataset format guide for COCO JSON and YOLO. Auto-detection, directory structure, annotation schemas, and format conversion.
+---
+
+# Dataset Formats
+
+RF-DETR supports training on datasets in two popular formats: **COCO** and **YOLO**. The format is automatically detected based on your dataset's directory structure—simply pass your dataset directory to the `train()` method.
+
+## Automatic Format Detection
+
+When you call `model.train(dataset_dir=<path>)`, RF-DETR checks the following:
+
+1. **COCO format**: Looks for `train/_annotations.coco.json`
+2. **YOLO format**: Looks for `data.yaml` (or `data.yml`) and `train/images/` directory
+
+If neither format is detected, an error is raised with instructions on what's expected.
+
+!!! tip "Roboflow Export"
+
+    [Roboflow](https://roboflow.com/annotate) can export datasets in both COCO and YOLO formats. When downloading from Roboflow, select the appropriate format based on your preference.
+
+---
+
+## COCO Format
+
+COCO (Common Objects in Context) format uses JSON files to store annotations in a structured format with images, categories, and annotations.
+
+### Directory Structure
+
+```
+dataset/
+├── train/
+│   ├── _annotations.coco.json
+│   ├── image1.jpg
+│   ├── image2.jpg
+│   └── ... (other image files)
+├── valid/
+│   ├── _annotations.coco.json
+│   ├── image1.jpg
+│   ├── image2.jpg
+│   └── ... (other image files)
+└── test/
+    ├── _annotations.coco.json
+    ├── image1.jpg
+    ├── image2.jpg
+    └── ... (other image files)
+```
+
+### Annotation File Structure
+
+Each `_annotations.coco.json` file contains:
+
+```json
+{
+  "info": {
+    "description": "Dataset description",
+    "version": "1.0"
+  },
+  "licenses": [],
+  "images": [
+    {
+      "id": 1,
+      "file_name": "image1.jpg",
+      "width": 640,
+      "height": 480
+    }
+  ],
+  "categories": [
+    {
+      "id": 1,
+      "name": "cat",
+      "supercategory": "animal"
+    },
+    {
+      "id": 2,
+      "name": "dog",
+      "supercategory": "animal"
+    }
+  ],
+  "annotations": [
+    {
+      "id": 1,
+      "image_id": 1,
+      "category_id": 1,
+      "bbox": [
+        100,
+        150,
+        200,
+        180
+      ],
+      "area": 36000,
+      "iscrowd": 0
+    }
+  ]
+}
+```
+
+#### Key Fields
+
+| Field         | Description                                                           |
+| ------------- | --------------------------------------------------------------------- |
+| `images`      | List of image metadata including `id`, `file_name`, `width`, `height` |
+| `categories`  | List of object categories with `id` and `name`                        |
+| `annotations` | List of object annotations linking images to categories               |
+| `bbox`        | Bounding box in `[x, y, width, height]` format (top-left corner)      |
+| `area`        | Area of the bounding box                                              |
+| `iscrowd`     | 0 for individual objects, 1 for crowd regions                         |
+
+### Segmentation Annotations
+
+For training segmentation models, your COCO annotations must include a `segmentation` key with polygon coordinates:
+
+```json
+{
+  "id": 1,
+  "image_id": 1,
+  "category_id": 1,
+  "bbox": [
+    100,
+    150,
+    200,
+    180
+  ],
+  "area": 36000,
+  "iscrowd": 0,
+  "segmentation": [
+    [
+      100,
+      150,
+      150,
+      150,
+      200,
+      200,
+      150,
+      250,
+      100,
+      200
+    ]
+  ]
+}
+```
+
+The `segmentation` field contains a list of polygons, where each polygon is a flat list of coordinates: `[x1, y1, x2, y2, x3, y3, ...]`.
+
+---
+
+### Keypoint Annotations
+
+For training the keypoint preview model, use COCO JSON keypoint annotations. Roboflow-style COCO exports are supported
+when the split files are named `train/_annotations.coco.json` and `valid/_annotations.coco.json`.
+
+Each keypoint annotation must include a bounding box plus COCO keypoint fields:
+
+```json
+{
+  "id": 1,
+  "image_id": 1,
+  "category_id": 0,
+  "bbox": [
+    100,
+    150,
+    200,
+    180
+  ],
+  "area": 36000,
+  "iscrowd": 0,
+  "num_keypoints": 17,
+  "keypoints": [
+    110,
+    160,
+    2,
+    125,
+    158,
+    2
+  ]
+}
+```
+
+The category should declare the keypoint schema:
+
+```json
+{
+  "id": 0,
+  "name": "person",
+  "supercategory": "person",
+  "keypoints": [
+    "nose",
+    "left_eye",
+    "right_eye"
+  ],
+  "skeleton": []
+}
+```
+
+The `keypoints` array above is shortened for readability. In a valid COCO person-keypoint annotation it contains
+`17 * 3` values: `x`, `y`, and visibility for each keypoint.
+
+The keypoint preview model is pretrained on COCO person-style keypoints. Its default COCO schema is `[17]`, so
+keypoint-bearing categories are mapped onto the active keypoint label slot during COCO loading. Legacy checkpoints may
+still report a background-first `[0, 17]` schema, which RF-DETR accepts for compatibility. Custom keypoint training can
+also use YOLO pose labels, described below.
+
+---
+
+## YOLO Format
+
+YOLO format uses separate text files for each image's annotations and a `data.yaml` configuration file that defines class names.
+
+### Directory Structure
+
+```
+dataset/
+├── data.yaml
+├── train/
+│   ├── images/
+│   │   ├── image1.jpg
+│   │   ├── image2.jpg
+│   │   └── ...
+│   └── labels/
+│       ├── image1.txt
+│       ├── image2.txt
+│       └── ...
+├── valid/
+│   ├── images/
+│   │   ├── image1.jpg
+│   │   ├── image2.jpg
+│   │   └── ...
+│   └── labels/
+│       ├── image1.txt
+│       ├── image2.txt
+│       └── ...
+└── test/
+    ├── images/
+    │   ├── image1.jpg
+    │   └── ...
+    └── labels/
+        ├── image1.txt
+        └── ...
+```
+
+### data.yaml Configuration
+
+The `data.yaml` file at the root of your dataset directory defines the class names:
+
+```yaml
+names:
+  - cat
+  - dog
+  - bird
+
+nc: 3
+
+train: train/images
+val: valid/images
+test: test/images
+```
+
+| Field                  | Description                                        |
+| ---------------------- | -------------------------------------------------- |
+| `names`                | List of class names (0-indexed)                    |
+| `nc`                   | Number of classes                                  |
+| `train`, `val`, `test` | Paths to image directories (relative to data.yaml) |
+
+!!! note "Alternative format"
+
+    Some YOLO datasets use a dictionary format for names:
+
+    ```yaml
+    names:
+      0: cat
+      1: dog
+      2: bird
+    ```
+
+    Both formats are supported.
+
+### Label File Format
+
+Each image has a corresponding `.txt` file in the `labels/` directory with the same base name. Each line in the label file represents one object:
+
+```
+<class_id> <x_center> <y_center> <width> <height>
+```
+
+**Example** (`image1.txt`):
+
+```
+0 0.5 0.4 0.3 0.2
+1 0.2 0.6 0.15 0.25
+```
+
+#### Coordinate Format
+
+| Field      | Range        | Description                                     |
+| ---------- | ------------ | ----------------------------------------------- |
+| `class_id` | 0, 1, 2, ... | Zero-indexed class ID from `names` in data.yaml |
+| `x_center` | 0.0 - 1.0    | Normalized x-coordinate of bounding box center  |
+| `y_center` | 0.0 - 1.0    | Normalized y-coordinate of bounding box center  |
+| `width`    | 0.0 - 1.0    | Normalized width of bounding box                |
+| `height`   | 0.0 - 1.0    | Normalized height of bounding box               |
+
+All coordinates are normalized relative to image dimensions. For example, if an image is 640×480 pixels and the bounding box center is at (320, 240):
+
+- `x_center` = 320 / 640 = 0.5
+- `y_center` = 240 / 480 = 0.5
+
+### Segmentation Labels (YOLO-Seg)
+
+For segmentation, YOLO format extends the label format with polygon coordinates:
+
+```
+<class_id> <x1> <y1> <x2> <y2> <x3> <y3> ...
+```
+
+**Example** (`image1.txt` with segmentation):
+
+```
+0 0.1 0.2 0.3 0.2 0.4 0.5 0.2 0.6 0.1 0.4
+```
+
+The coordinates after the class ID represent the polygon vertices in normalized format.
+
+---
+
+### Pose Labels (YOLO Pose)
+
+For keypoint preview training, RF-DETR supports Ultralytics YOLO pose labels in the same directory layout shown above.
+The `data.yaml` file must declare `kpt_shape`:
+
+```yaml
+names:
+  0: person
+
+kpt_shape: [17, 3] # [number_of_keypoints, dimensions]; dimensions must be 2 or 3
+flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+kpt_names:
+  0:
+    - nose
+    - left_eye
+    - right_eye
+```
+
+`kpt_names` is optional. When omitted, RF-DETR creates placeholder names such as `keypoint_0`. `flip_idx` is an
+Ultralytics-style length-`K` permutation used to infer RF-DETR's flat `keypoint_flip_pairs` for horizontal-flip
+augmentation.
+
+Each pose label row contains a bounding box followed by keypoints:
+
+```text
+<class_id> <x_center> <y_center> <width> <height> <px1> <py1> <v1> ... <pxK> <pyK> <vK>
+```
+
+For `kpt_shape: [K, 2]`, omit the visibility value:
+
+```text
+<class_id> <x_center> <y_center> <width> <height> <px1> <py1> ... <pxK> <pyK>
+```
+
+All box and keypoint coordinates are normalized to `[0, 1]`. RF-DETR converts keypoints to COCO-style `(x, y, visibility)` tensors internally. For `[K, 3]`, the visibility values are preserved. For `[K, 2]`, visibility is
+synthesized: nonzero points are marked visible (`2`) and `(0, 0)` points are marked absent (`0`).
+
+Use the YOLO schema helper when you want to configure a model explicitly:
+
+```python
+from pathlib import Path
+
+from rfdetr import RFDETRKeypointPreview
+from rfdetr.datasets._keypoint_schema import infer_yolo_keypoint_schema
+
+DATASET_DIR = Path("/path/to/yolo-pose-dataset")
+schema = infer_yolo_keypoint_schema(DATASET_DIR / "data.yaml")
+
+model = RFDETRKeypointPreview(
+    num_classes=len(schema.class_names),
+    num_keypoints_per_class=schema.num_keypoints_per_class,
+)
+
+model.train(
+    dataset_file="yolo",
+    dataset_dir=str(DATASET_DIR),
+    class_names=schema.class_names,
+    keypoint_oks_sigmas=schema.keypoint_oks_sigmas,
+)
+```
+
+!!! note "flip_idx and keypoint_flip_pairs"
+
+    `flip_idx` is a permutation, while `keypoint_flip_pairs` is a flat pair list. During `model.train()`, RF-DETR infers
+    the pair list automatically from `flip_idx` when no explicit `keypoint_flip_pairs` is provided.
+
+---
+
+## Converting Between Formats
+
+### YOLO to COCO
+
+You can use the [supervision](https://github.com/roboflow/supervision) library to convert datasets:
+
+```python
+import supervision as sv
+
+# Load YOLO dataset
+dataset = sv.DetectionDataset.from_yolo(
+    images_directory_path="path/to/images",
+    annotations_directory_path="path/to/labels",
+    data_yaml_path="path/to/data.yaml",
+)
+
+# Save as COCO
+dataset.as_coco(images_directory_path="output/images", annotations_path="output/annotations.json")
+```
+
+### COCO to YOLO
+
+```python
+import supervision as sv
+
+# Load COCO dataset
+dataset = sv.DetectionDataset.from_coco(
+    images_directory_path="path/to/images", annotations_path="path/to/annotations.json"
+)
+
+# Save as YOLO
+dataset.as_yolo(
+    images_directory_path="output/images", annotations_directory_path="output/labels", data_yaml_path="output/data.yaml"
+)
+```
+
+### Using Roboflow
+
+[Roboflow](https://roboflow.com) provides a web interface to:
+
+1. Upload datasets in any format
+2. Annotate new images or edit existing annotations
+3. Export in COCO, YOLO, or other formats
+
+This is often the easiest way to convert between formats while also having the option to augment your data.
+
+---
+
+## Which Format Should I Use?
+
+Both formats work equally well with RF-DETR. Choose based on your workflow:
+
+| Consideration                     | COCO                       | YOLO                    |
+| --------------------------------- | -------------------------- | ----------------------- |
+| **Annotation storage**            | Single JSON file per split | One text file per image |
+| **Human readability**             | JSON structure, verbose    | Simple text, compact    |
+| **Other framework compatibility** | DETR family, MMDetection   | Ultralytics YOLO        |
+| **Segmentation support**          | Full polygon support       | Full polygon support    |
+| **Editing annotations**           | Requires JSON parsing      | Simple text editing     |
+
+!!! tip "Recommendation"
+
+    If you're exporting from Roboflow or already have a dataset in one format, simply use that format. RF-DETR handles both identically.
+
+---
+
+## Troubleshooting
+
+### Format Detection Fails
+
+If you see an error like:
+
+```
+Could not detect dataset format in /path/to/dataset
+```
+
+Check that:
+
+**For COCO format:**
+
+- `train/_annotations.coco.json` exists
+- The JSON file is valid
+
+**For YOLO format:**
+
+- `data.yaml` or `data.yml` exists at the root
+- `train/images/` directory exists with images
+
+### Empty Annotations
+
+If images have no objects, handle them as follows:
+
+**COCO format:** Include the image in the `images` array but don't add any annotations for it.
+
+**YOLO format:** Create an empty `.txt` file (0 bytes) for the image, or omit the label file entirely.
+
+### Class ID Mismatch
+
+**COCO format:** Category IDs in annotations must match IDs defined in the `categories` array.
+
+**YOLO format:** Class IDs in label files must be valid indices (0 to `nc-1`) based on the `names` list in `data.yaml`.
@@ -0,0 +1,247 @@
+---
+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=4,
+        grad_accum_steps=4,
+        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=4,
+        grad_accum_steps=4,
+        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=2,
+        grad_accum_steps=8,
+        lr=1e-5,
+        output_dir="<OUTPUT_PATH>",
+    )
+    ```
+
+Different GPUs have different VRAM capacities, so adjust batch_size and grad_accum_steps to maintain a total batch size of 16. For example, on a powerful GPU like the A100, use `batch_size=16` and `grad_accum_steps=1`; on smaller GPUs like the T4, use `batch_size=4` and `grad_accum_steps=4`. This gradient accumulation strategy helps train effectively even with limited memory.
+
+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` + `train/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
+- [Custom augmentations with Albumentations](augmentations.md) - Dedicated guide
+- [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:
+
+- `checkpoint.pth` – the most recent checkpoint, saved at the end of the latest epoch.
+
+- `checkpoint_<number>.pth` – periodic checkpoints saved every N epochs (default is every 10).
+
+- `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 only the model weights (no optimizer state or scheduler) and is chosen as the better of the EMA and non-EMA models based on validation performance.
+
+For detection and segmentation models, the validation score is box mAP (`val/mAP_50_95`). 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.
+
+??? note "Checkpoint file sizes"
+
+    Checkpoint sizes vary based on what they contain:
+
+    - **Training checkpoints** (e.g. `checkpoint.pth`, `checkpoint_<number>.pth`) include model weights, optimizer state, scheduler state, and training metadata. Use these to resume training.
+
+    - **Evaluation checkpoints** (e.g. `checkpoint_best_ema.pth`, `checkpoint_best_regular.pth`) store only the model weights — either EMA or raw — and are used to track the best-performing models. These may come from different epochs depending on which version achieved the highest validation score.
+
+    - **Stripped checkpoint** (e.g. `checkpoint_best_total.pth`) contains only the final model weights and is optimized for inference and deployment.
+
+## Load and Run Fine-Tuned Model
+
+=== "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>")
+    ```
+
+## Next Steps
+
+After training your model, you can:
+
+- [Export your model to ONNX](../export.md) for deployment with various inference frameworks
+- [Deploy to Roboflow](../deploy.md) for cloud-based inference and workflow integration
@@ -0,0 +1,316 @@
+---
+description: Track RF-DETR training with TensorBoard, Weights and Biases, and MLflow. Configure multiple experiment loggers simultaneously.
+---
+
+# Training Loggers
+
+RF-DETR supports integration with popular experiment tracking and visualization platforms. You can enable one or more supported loggers to monitor your training runs, compare experiments, and track metrics over time.
+
+## CSV (always active)
+
+A `CSVLogger` is always active regardless of any flags. It requires no extra packages and writes all metrics to `{output_dir}/metrics.csv` on every validation step.
+
+---
+
+## TensorBoard
+
+[TensorBoard](https://www.tensorflow.org/tensorboard) is a powerful toolkit for visualizing and tracking training metrics.
+
+TensorBoard logging is enabled by default. Pass `tensorboard=False` to disable it.
+
+!!! note "Missing package behaviour"
+
+    If the `tensorboard` package is not installed, training continues without error — a
+    `UserWarning` is emitted and TensorBoard logging is silently suppressed. Install
+    `rfdetr[loggers]` to avoid this.
+
+### Setup
+
+Install the required packages:
+
+```bash
+pip install "rfdetr[loggers]"
+```
+
+### Usage
+
+TensorBoard is active unless you explicitly disable it:
+
+```python
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium()
+
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,
+    grad_accum_steps=4,
+    lr=1e-4,
+    output_dir="output",
+    # tensorboard=True is the default; pass tensorboard=False to disable
+)
+```
+
+### Viewing Logs
+
+**Local environment:**
+
+```bash
+tensorboard --logdir output
+```
+
+Then open `http://localhost:6006/` in your browser.
+
+**Google Colab:**
+
+```ipython
+%load_ext tensorboard
+%tensorboard --logdir output
+```
+
+### Logged Metrics
+
+All logged metric keys are listed in the [Logged Metrics Reference](customization.md#logged-metrics-reference).
+
+---
+
+## Weights and Biases
+
+[Weights and Biases (W&B)](https://www.wandb.ai) is a cloud-based platform for experiment tracking and visualization.
+
+### Setup
+
+Install the required packages:
+
+```bash
+pip install "rfdetr[loggers]"
+```
+
+Log in to W&B:
+
+```bash
+wandb login
+```
+
+You can retrieve your API key at [wandb.ai/authorize](https://wandb.ai/authorize).
+
+### Usage
+
+Enable W&B logging in your training:
+
+```python
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium()
+
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,
+    grad_accum_steps=4,
+    lr=1e-4,
+    output_dir="output",
+    wandb=True,
+    project="my-detection-project",
+    run="experiment-001",
+)
+```
+
+### Configuration
+
+| Parameter | Description                             |
+| --------- | --------------------------------------- |
+| `project` | Groups related experiments together     |
+| `run`     | Identifies individual training sessions |
+
+If you don't specify a run name, W&B assigns a random one automatically.
+
+### Features
+
+Access your runs at [wandb.ai](https://wandb.ai). W&B provides:
+
+- Real-time metric visualization
+- Experiment comparison
+- Hyperparameter tracking
+- System metrics (GPU usage, memory)
+- Training config logging
+
+### Logged Metrics
+
+All logged metric keys are listed in the [Logged Metrics Reference](customization.md#logged-metrics-reference).
+
+---
+
+## ClearML
+
+[ClearML](https://clear.ml) is an open-source platform for managing, tracking, and automating machine learning experiments.
+
+**ClearML is not yet integrated as a native PTL logger.** Passing `clearml=True` to `model.train()` raises `NotImplementedError`; metrics are not logged to ClearML through RF-DETR's built-in logger wiring.
+
+### Workaround: ClearML SDK auto-binding
+
+ClearML's SDK captures PyTorch Lightning metrics automatically when a `Task` is initialised before training begins:
+
+```python
+from clearml import Task
+from rfdetr import RFDETRMedium
+
+# Initialise before model.train() — ClearML auto-binds to PTL logging
+task = Task.init(project_name="my-detection-project", task_name="experiment-001")
+
+model = RFDETRMedium()
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,
+    grad_accum_steps=4,
+    lr=1e-4,
+    output_dir="output",
+    # Do NOT pass clearml=True — RF-DETR raises NotImplementedError for that flag
+)
+```
+
+Alternatively, attach a ClearML callback directly using the [Custom Training API](#attaching-loggers-via-the-custom-training-api).
+
+---
+
+## MLflow
+
+[MLflow](https://mlflow.org/) is an open-source platform for the machine learning lifecycle that helps track experiments, package code into reproducible runs, and share and deploy models.
+
+### Setup
+
+Install the required packages:
+
+```bash
+pip install "rfdetr[loggers]"
+```
+
+### Usage
+
+Enable MLflow logging in your training:
+
+```python
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium()
+
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,
+    grad_accum_steps=4,
+    lr=1e-4,
+    output_dir="output",
+    mlflow=True,
+    project="my-detection-project",
+    run="experiment-001",
+)
+```
+
+### Configuration
+
+| Parameter | Description                                         |
+| --------- | --------------------------------------------------- |
+| `project` | Sets the experiment name in MLflow                  |
+| `run`     | Sets the run name (auto-generated if not specified) |
+
+### Custom Tracking Server
+
+To use a custom MLflow tracking server, set environment variables:
+
+```python
+import os
+
+# Set MLflow tracking URI
+os.environ["MLFLOW_TRACKING_URI"] = "https://your-mlflow-server.com"
+
+# For authentication with tracking servers that require it
+os.environ["MLFLOW_TRACKING_TOKEN"] = "your-auth-token"
+
+# Then initialize and train your model
+model = RFDETRMedium()
+model.train(..., mlflow=True)
+```
+
+For teams using a hosted MLflow service (like Databricks), you'll typically need to set:
+
+- `MLFLOW_TRACKING_URI`: The URL of your MLflow tracking server
+- `MLFLOW_TRACKING_TOKEN`: Authentication token for your MLflow server
+
+### Viewing Logs
+
+Start the MLflow UI:
+
+```bash
+mlflow ui --backend-store-uri <OUTPUT_PATH>
+```
+
+Then open `http://localhost:5000` in your browser to access the MLflow dashboard.
+
+### Logged Metrics
+
+All logged metric keys are listed in the [Logged Metrics Reference](customization.md#logged-metrics-reference).
+
+---
+
+## Using Multiple Loggers
+
+You can enable multiple logging systems simultaneously:
+
+```python
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    tensorboard=True,
+    wandb=True,
+    mlflow=True,
+    project="my-project",
+    run="experiment-001",
+)
+```
+
+This allows you to leverage the strengths of different platforms:
+
+- **TensorBoard**: Local visualization and debugging
+- **W&B**: Cloud-based collaboration and experiment comparison
+- **MLflow**: Model registry and deployment tracking
+
+Note: `clearml=True` is accepted by the config schema but raises `NotImplementedError` when the trainer is built. Use the [ClearML SDK workaround](#clearml) instead.
+
+---
+
+## Attaching loggers via the Custom Training API
+
+`build_trainer` automatically creates loggers from `TrainConfig` flags. To attach a logger not listed above (for example Neptune, Comet, or a fully custom logger), build it separately and append it to `trainer.loggers` before calling `trainer.fit`:
+
+```python
+from rfdetr.config import RFDETRMediumConfig, TrainConfig
+from rfdetr.training import RFDETRModelModule, RFDETRDataModule, build_trainer
+
+model_config = RFDETRMediumConfig(num_classes=10)
+train_config = TrainConfig(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    output_dir="output",
+    tensorboard=True,  # built-in loggers still work
+)
+
+module = RFDETRModelModule(model_config, train_config)
+datamodule = RFDETRDataModule(model_config, train_config)
+trainer = build_trainer(train_config, model_config)
+
+# Attach any additional PTL-compatible logger
+from pytorch_lightning.loggers import CSVLogger  # example — use any PTL logger
+
+trainer.loggers.append(CSVLogger(save_dir="output", name="extra"))
+
+trainer.fit(module, datamodule)
+```
+
+CSVLogger is always active (it requires no extra packages). All logged metric keys — `train/loss`, `val/mAP_50_95`,
+`val/keypoint_map_50_95`, `val/F1`, `val/ema_mAP_50_95`, `val/AP/<class>`, etc. — are written to every logger in the
+list.
+
+→ **[Full list of logged metrics](customization.md#logged-metrics-reference)**
@@ -0,0 +1,288 @@
+---
+description: Complete RF-DETR training parameter reference. Learning rate, batch size, EMA, early stopping, resolution, and hardware configuration.
+---
+
+# Training Parameters
+
+This page provides a complete reference of all parameters available when training RF-DETR models.
+
+## Basic Example
+
+```python
+from rfdetr import RFDETRMedium
+
+model = RFDETRMedium()
+
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    batch_size=4,
+    grad_accum_steps=4,
+    lr=1e-4,
+    output_dir="output",
+)
+```
+
+## Core Parameters
+
+These are the essential parameters for training:
+
+| Parameter          | Type            | Default    | Description                                                                                                                                                       |
+| ------------------ | --------------- | ---------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `dataset_dir`      | `str`           | Required   | Path to your dataset directory. RF-DETR auto-detects if it's in COCO or YOLO format. See [Dataset Formats](dataset-formats.md).                                   |
+| `output_dir`       | `str`           | `"output"` | Directory where training artifacts (checkpoints, logs) are saved.                                                                                                 |
+| `epochs`           | `int`           | `100`      | Number of full passes over the training dataset.                                                                                                                  |
+| `batch_size`       | `int or "auto"` | `4`        | Number of samples processed per iteration. Higher values require more GPU memory. Set to `"auto"` to probe the GPU for the largest safe batch size automatically. |
+| `grad_accum_steps` | `int`           | `4`        | Accumulates gradients over multiple mini-batches. Use with `batch_size` to achieve effective batch size.                                                          |
+| `resume`           | `str`           | `None`     | Path to a saved checkpoint to continue training. Restores model weights, optimizer state, and scheduler.                                                          |
+
+### Understanding Batch Size
+
+The **effective batch size** is calculated as:
+
+```
+effective_batch_size = batch_size × grad_accum_steps × num_gpus
+```
+
+Recommended configurations for different GPUs (targeting effective batch size of 16):
+
+| GPU      | VRAM    | `batch_size` | `grad_accum_steps` |
+| -------- | ------- | ------------ | ------------------ |
+| A100     | 40-80GB | 16           | 1                  |
+| RTX 4090 | 24GB    | 8            | 2                  |
+| RTX 3090 | 24GB    | 8            | 2                  |
+| T4       | 16GB    | 4            | 4                  |
+| RTX 3070 | 8GB     | 2            | 8                  |
+
+## Learning Rate Parameters
+
+| Parameter    | Type    | Default  | Description                                                                                                                                                          |
+| ------------ | ------- | -------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `lr`         | `float` | `1e-4`   | Learning rate for most parts of the model.                                                                                                                           |
+| `lr_encoder` | `float` | `1.5e-4` | Learning rate specifically for the backbone encoder. Can be set lower than `lr` if you want to fine-tune the encoder more conservatively than the rest of the model. |
+
+!!! tip "Learning rate tips"
+
+    - Start with the default values for fine-tuning
+    - If the model doesn't converge, try reducing `lr` by half
+    - For training from scratch (not recommended), you may need higher learning rates
+
+## Resolution Parameters
+
+| Parameter    | Type  | Default         | Description                                                                                                                                                                                                                                                                                                                                                                                                   |
+| ------------ | ----- | --------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `resolution` | `int` | Model-dependent | Input image resolution. Higher values can improve accuracy but require more memory. Each model has its own valid block size: current standard detection checkpoints use multiples of 32, current segmentation checkpoints use multiples of 24 (most variants) or 12 (`RFDETRSegNano`), and the definitive rule is that the resolution must be divisible by `patch_size * num_windows` for the selected model. |
+
+Common resolution values for currently documented checkpoints:
+
+- Detection: `384`, `512`, `576`, `704`
+- Segmentation: `312`, `384`, `432`, `504`, `624`, `768`
+
+For example, `RFDETRSegXLarge` uses `624x624`, which is valid because `624` is divisible by `24`.
+
+## Regularization Parameters
+
+| Parameter      | Type    | Default | Description                                                                           |
+| -------------- | ------- | ------- | ------------------------------------------------------------------------------------- |
+| `weight_decay` | `float` | `1e-4`  | L2 regularization coefficient. Helps prevent overfitting by penalizing large weights. |
+
+## Hardware Parameters
+
+| Parameter                | Type   | Default | Description                                                                                                                                                                                                                                    |
+| ------------------------ | ------ | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `device`                 | `str`  | `None`  | Device to run training on. `None` means auto-detected by PyTorch Lightning. Options: `"cuda"`, `"cpu"`, `"mps"` (Apple Silicon).                                                                                                               |
+| `gradient_checkpointing` | `bool` | `False` | **Constructor-only parameter** — pass to the model constructor (`RFDETRMedium(gradient_checkpointing=True)`), not to `train()`. Re-computes activations during backprop to reduce memory usage by ~30-40% at the cost of ~20% slower training. |
+
+## EMA (Exponential Moving Average)
+
+| Parameter | Type   | Default | Description                                                                                                          |
+| --------- | ------ | ------- | -------------------------------------------------------------------------------------------------------------------- |
+| `use_ema` | `bool` | `True`  | Enables Exponential Moving Average of weights. Produces a smoothed checkpoint that often improves final performance. |
+
+!!! info "What is EMA?"
+
+    EMA maintains a moving average of the model weights throughout training. This smoothed version often generalizes better than the raw weights and is commonly used for the final model.
+
+## Checkpoint Parameters
+
+| Parameter             | Type  | Default | Description                                                                                                                            |
+| --------------------- | ----- | ------- | -------------------------------------------------------------------------------------------------------------------------------------- |
+| `checkpoint_interval` | `int` | `10`    | Frequency (in epochs) at which model checkpoints are saved. More frequent saves provide better coverage but consume more storage.      |
+| `skip_best_epochs`    | `int` | `0`     | Ignore the first N epochs when tracking best checkpoints and early-stopping patience. Useful when fine-tuning from a prior checkpoint. |
+
+### Checkpoint Files
+
+During training, multiple checkpoints are saved:
+
+| File                          | Description                               |
+| ----------------------------- | ----------------------------------------- |
+| `checkpoint.pth`              | Most recent checkpoint (for resuming)     |
+| `checkpoint_<N>.pth`          | Periodic checkpoint at epoch N            |
+| `checkpoint_best_ema.pth`     | Best validation performance (EMA weights) |
+| `checkpoint_best_regular.pth` | Best validation performance (raw weights) |
+| `checkpoint_best_total.pth`   | Final best model for inference            |
+
+Best validation performance uses the task metric for the model family: box mAP for detection/segmentation and COCO
+keypoint AP for keypoint preview.
+
+## Early Stopping Parameters
+
+| Parameter                  | Type    | Default | Description                                                                              |
+| -------------------------- | ------- | ------- | ---------------------------------------------------------------------------------------- |
+| `early_stopping`           | `bool`  | `False` | Enable early stopping based on the validation task metric.                               |
+| `early_stopping_patience`  | `int`   | `10`    | Number of epochs without improvement before stopping.                                    |
+| `early_stopping_min_delta` | `float` | `0.001` | Minimum metric change to qualify as an improvement.                                      |
+| `early_stopping_use_ema`   | `bool`  | `False` | Whether to track improvements using EMA model metrics.                                   |
+| `skip_best_epochs`         | `int`   | `0`     | Ignore the first N epochs (0..N-1) for best-model selection and early-stopping patience. |
+
+### Early Stopping Example
+
+```python
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=200,
+    batch_size=4,
+    early_stopping=True,
+    early_stopping_patience=15,
+    early_stopping_min_delta=0.005,
+    skip_best_epochs=3,
+)
+```
+
+This configuration will:
+
+- Train for up to 200 epochs
+- Ignore epochs 0-2 for best-checkpoint tracking and patience counting
+- Stop early if the validation metric doesn't improve by at least 0.005 for 15 consecutive epochs
+
+!!! note "Transfer learning with `pretrain_weights`"
+
+    When fine-tuning from `pretrain_weights`, the pretrained model's epoch-0 validation metric can be artificially high
+    relative to the training trajectory on the new dataset. This causes `checkpoint_best_total.pth` to always contain
+    the untrained pretrained weights and may trigger early stopping prematurely. Use `skip_best_epochs` to defer
+    best-checkpoint selection and patience counting until the model has had time to adapt.
+
+## Logging Parameters
+
+| Parameter     | Type   | Default | Description                                                                                                                                                                                                 |
+| ------------- | ------ | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `tensorboard` | `bool` | `True`  | Enable TensorBoard logging. Requires `pip install "rfdetr[loggers]"`. If the `tensorboard` package is not installed, training continues with a `UserWarning` and TensorBoard output is silently suppressed. |
+| `wandb`       | `bool` | `False` | Enable Weights & Biases logging. Requires `pip install "rfdetr[loggers]"`.                                                                                                                                  |
+| `project`     | `str`  | `None`  | Project name for W&B logging.                                                                                                                                                                               |
+| `run`         | `str`  | `None`  | Run name for W&B logging. If not specified, W&B assigns a random name.                                                                                                                                      |
+
+### Logging Example
+
+```python
+model.train(
+    dataset_dir="path/to/dataset",
+    epochs=100,
+    tensorboard=True,
+    wandb=True,
+    project="my-detection-project",
+    run="experiment-001",
+)
+```
+
+## Evaluation Parameters
+
+| Parameter               | Type                | Default | Description                                                                                                        |
+| ----------------------- | ------------------- | ------- | ------------------------------------------------------------------------------------------------------------------ |
+| `eval_max_dets`         | `int`               | `500`   | Maximum number of detections per image considered during COCO evaluation. Lower values speed up evaluation.        |
+| `eval_interval`         | `int`               | `1`     | Run COCO evaluation every N epochs. Set to a higher value to reduce evaluation overhead during long training runs. |
+| `log_per_class_metrics` | `bool`              | `True`  | Log per-class AP metrics to the console and loggers. Disable to reduce log verbosity when there are many classes.  |
+| `progress_bar`          | str \| bool \| None | `None`  | Progress bar style: `"tqdm"`, `"rich"`, or `None`. Legacy booleans are still accepted.                             |
+
+## Keypoint Preview Parameters
+
+These parameters apply when training `RFDETRKeypointPreview` on COCO keypoint annotations or Ultralytics YOLO pose labels.
+
+| Parameter                     | Type                  | Default | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               |
+| ----------------------------- | --------------------- | ------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `num_keypoints_per_class`     | `list[int]`           | `[17]`  | **Constructor parameter** — pass to `RFDETRKeypointPreview(num_keypoints_per_class=...)`. Keypoint schema by model label slot. A zero entry marks a detection-only class slot; legacy checkpoints may use a background-first `[0, 17]` schema.                                                                                                                                                                                                                                                            |
+| `keypoint_flip_pairs`         | `list[int]`           | `[]`    | Flat left/right keypoint index pairs used to swap joints after horizontal-flip augmentation. YOLO `flip_idx` metadata is a permutation; RF-DETR converts it to this pair-list form during automatic schema inference when possible — it extracts only symmetric mutual pairs where `flip_idx[i] == j` and `flip_idx[j] == i`. Asymmetric entries and self-mapped keypoints (`flip_idx[i] == i`) are silently omitted; supply `keypoint_flip_pairs` explicitly when your `flip_idx` includes such entries. |
+| `keypoint_l1_loss_coef`       | `float`               | `1.0`   | Weight for keypoint coordinate L1 loss in keypoint preview training.                                                                                                                                                                                                                                                                                                                                                                                                                                      |
+| `keypoint_findable_loss_coef` | `float`               | `1.0`   | Weight for keypoint findable/objectness loss.                                                                                                                                                                                                                                                                                                                                                                                                                                                             |
+| `keypoint_visible_loss_coef`  | `float`               | `1.0`   | Weight for keypoint visibility loss.                                                                                                                                                                                                                                                                                                                                                                                                                                                                      |
+| `keypoint_nll_loss_coef`      | `float`               | `1.0`   | Weight for keypoint negative-log-likelihood loss.                                                                                                                                                                                                                                                                                                                                                                                                                                                         |
+| `keypoint_oks_sigmas`         | `list[float] \| None` | `None`  | Per-keypoint OKS sigma values used for COCO AP evaluation. When `None`, 17-keypoint person datasets use the evaluator's standard COCO sigmas and custom keypoint counts use RF-DETR's uniform custom fallback. Pass explicit values, such as schema-inferred sigmas, when you need a specific custom OKS policy.                                                                                                                                                                                          |
+
+!!! note "OKS sigma values: flat vs per-keypoint"
+
+    `infer_coco_keypoint_schema` and `infer_yolo_keypoint_schema` return a flat sigma of 0.1 for all inferred keypoints, and the keypoint demos pass those values explicitly for custom datasets. If `keypoint_oks_sigmas=None`, COCO person-keypoint evaluation uses the standard 17-keypoint COCO sigmas, while non-17 custom keypoint counts use RF-DETR's uniform custom fallback. Flat custom sigmas are not directly comparable to official COCO benchmark numbers.
+
+## Advanced Parameters
+
+The parameters below are available for fine-grained control over training behaviour. Most users can leave these at their defaults.
+
+### Scheduler and Regularization
+
+| Parameter       | Type    | Default  | Description                                                                                                 |
+| --------------- | ------- | -------- | ----------------------------------------------------------------------------------------------------------- |
+| `lr_scheduler`  | `str`   | `"step"` | Learning rate scheduler type. Options: `"step"` (step decay at `lr_drop`) or `"cosine"` (cosine annealing). |
+| `lr_min_factor` | `float` | `0.0`    | Floor for the cosine scheduler, expressed as a fraction of the initial LR. Ignored when using `"step"`.     |
+| `warmup_epochs` | `float` | `0.0`    | Number of epochs for linear learning rate warmup at the start of training.                                  |
+| `drop_path`     | `float` | `0.0`    | Stochastic depth drop-path rate applied to the backbone. Higher values add more regularization.             |
+
+### Runtime and Accelerator
+
+| Parameter           | Type   | Default  | Description                                                                                      |
+| ------------------- | ------ | -------- | ------------------------------------------------------------------------------------------------ |
+| `accelerator`       | `str`  | `"auto"` | PyTorch Lightning accelerator selection. `"auto"` picks GPU if available, then MPS, then CPU.    |
+| `seed`              | `int`  | `None`   | Global random seed for reproducibility. `None` means no fixed seed is set.                       |
+| `fp16_eval`         | `bool` | `False`  | Run evaluation passes in FP16 precision. Reduces memory usage but may lower numerical precision. |
+| `compute_val_loss`  | `bool` | `True`   | Compute and log the detection loss on the validation set each epoch.                             |
+| `compute_test_loss` | `bool` | `True`   | Compute and log the detection loss during the final test run.                                    |
+
+### DataLoader Tuning
+
+| Parameter            | Type   | Default | Description                                                                                               |
+| -------------------- | ------ | ------- | --------------------------------------------------------------------------------------------------------- |
+| `pin_memory`         | `bool` | `None`  | Pin host memory in the DataLoader for faster GPU transfers. `None` defers to PyTorch Lightning's default. |
+| `persistent_workers` | `bool` | `None`  | Keep DataLoader worker processes alive between epochs. `None` defers to PyTorch Lightning's default.      |
+| `prefetch_factor`    | `int`  | `None`  | Number of batches to prefetch per DataLoader worker. `None` uses PyTorch's built-in default.              |
+
+## Complete Parameter Reference
+
+Below is a summary table of all training parameters:
+
+| Parameter                  | Type                | Default        | Description                                                                                                                           |
+| -------------------------- | ------------------- | -------------- | ------------------------------------------------------------------------------------------------------------------------------------- |
+| `dataset_dir`              | str                 | Required       | Path to COCO or YOLO formatted dataset with train/valid/test splits.                                                                  |
+| `output_dir`               | str                 | "output"       | Directory for checkpoints, logs, and other training artifacts.                                                                        |
+| `epochs`                   | int                 | 100            | Number of full passes over the dataset.                                                                                               |
+| `batch_size`               | int or "auto"       | 4              | Samples per iteration. Set to `"auto"` to let RF-DETR probe the GPU for the largest safe batch size. Balance with `grad_accum_steps`. |
+| `grad_accum_steps`         | int                 | 4              | Gradient accumulation steps for effective larger batch sizes.                                                                         |
+| `lr`                       | float               | 1e-4           | Learning rate for the model (excluding encoder).                                                                                      |
+| `lr_encoder`               | float               | 1.5e-4         | Learning rate for the backbone encoder.                                                                                               |
+| `resolution`               | int                 | Model-specific | Input image size (must be divisible by the selected model's `patch_size * num_windows`).                                              |
+| `weight_decay`             | float               | 1e-4           | L2 regularization coefficient.                                                                                                        |
+| `device`                   | str                 | "cuda"         | Training device: cuda, cpu, or mps.                                                                                                   |
+| `use_ema`                  | bool                | True           | Enable Exponential Moving Average of weights.                                                                                         |
+| `gradient_checkpointing`   | bool                | False          | Trade compute for memory during backprop.                                                                                             |
+| `checkpoint_interval`      | int                 | 10             | Save checkpoint every N epochs.                                                                                                       |
+| `resume`                   | str                 | None           | Path to checkpoint for resuming training.                                                                                             |
+| `tensorboard`              | bool                | True           | Enable TensorBoard logging.                                                                                                           |
+| `wandb`                    | bool                | False          | Enable Weights & Biases logging.                                                                                                      |
+| `project`                  | str                 | None           | W&B project name.                                                                                                                     |
+| `run`                      | str                 | None           | W&B run name.                                                                                                                         |
+| `early_stopping`           | bool                | False          | Enable early stopping.                                                                                                                |
+| `early_stopping_patience`  | int                 | 10             | Epochs without improvement before stopping.                                                                                           |
+| `early_stopping_min_delta` | float               | 0.001          | Minimum validation metric change to qualify as improvement.                                                                           |
+| `early_stopping_use_ema`   | bool                | False          | Use EMA model for early stopping metrics.                                                                                             |
+| `eval_max_dets`            | int                 | 500            | Maximum detections per image considered during COCO evaluation.                                                                       |
+| `eval_interval`            | int                 | 1              | Run COCO evaluation every N epochs.                                                                                                   |
+| `log_per_class_metrics`    | bool                | True           | Log per-class AP metrics to the console and loggers.                                                                                  |
+| `progress_bar`             | str \| bool \| None | None           | Progress bar style: `"tqdm"`, `"rich"`, or `None`. Legacy booleans are still accepted.                                                |
+| `accelerator`              | str                 | "auto"         | PyTorch Lightning accelerator. "auto" selects GPU/MPS/CPU automatically.                                                              |
+| `seed`                     | int                 | None           | Random seed for reproducibility. None means no fixed seed.                                                                            |
+| `lr_scheduler`             | str                 | "step"         | Learning rate scheduler type: "step" or "cosine".                                                                                     |
+| `lr_min_factor`            | float               | 0.0            | Minimum LR as a fraction of the initial LR (cosine scheduler floor).                                                                  |
+| `warmup_epochs`            | float               | 0.0            | Number of linear warmup epochs at the start of training.                                                                              |
+| `drop_path`                | float               | 0.0            | Stochastic depth drop-path rate for the backbone.                                                                                     |
+| `compute_val_loss`         | bool                | True           | Compute and log loss during validation.                                                                                               |
+| `compute_test_loss`        | bool                | True           | Compute and log loss during the test run.                                                                                             |
+| `fp16_eval`                | bool                | False          | Run evaluation in FP16 precision to reduce memory usage.                                                                              |
+| `pin_memory`               | bool                | None           | Pin DataLoader memory. None defers to PyTorch Lightning's default.                                                                    |
+| `persistent_workers`       | bool                | None           | Keep DataLoader workers alive between epochs. None uses PTL default.                                                                  |
+| `prefetch_factor`          | int                 | None           | Number of batches prefetched per worker. None uses PyTorch default.                                                                   |