chore: import upstream snapshot with attribution
CPU tests Workflow / Testing (ubuntu-latest, 3.12) (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.13) (push) Failing after 0s
Mypy Type Check / Type Check (push) Failing after 0s
Docs/Test WorkFlow / Test docs build (push) Failing after 1s
PR Conflict Labeler / labeling (push) Failing after 1s
Dependency resolution / Resolve [tflite] extra — Python 3.12 (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.10) (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.13) (push) Failing after 1s
CPU tests Workflow / build-pkg (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.10) (push) Failing after 0s
CPU tests Workflow / Testing (ubuntu-latest, 3.11) (push) Failing after 0s
Smoke Tests / try-all-models (macos-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (macos-latest, 3.13) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / testing-guardian (push) Has been cancelled
GPU tests Workflow / Testing (push) Has been cancelled
CPU tests Workflow / Testing (ubuntu-latest, 3.12) (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.13) (push) Failing after 0s
Mypy Type Check / Type Check (push) Failing after 0s
Docs/Test WorkFlow / Test docs build (push) Failing after 1s
PR Conflict Labeler / labeling (push) Failing after 1s
Dependency resolution / Resolve [tflite] extra — Python 3.12 (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.10) (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.13) (push) Failing after 1s
CPU tests Workflow / build-pkg (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.10) (push) Failing after 0s
CPU tests Workflow / Testing (ubuntu-latest, 3.11) (push) Failing after 0s
Smoke Tests / try-all-models (macos-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (macos-latest, 3.13) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / testing-guardian (push) Has been cancelled
GPU tests Workflow / Testing (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,112 @@
|
||||
---
|
||||
description: RF-DETR benchmark results on COCO and RF100-VL for detection, segmentation, and keypoint detection. Compare accuracy and latency against YOLO, D-FINE, and LW-DETR.
|
||||
---
|
||||
|
||||
# Benchmarks
|
||||
|
||||
!!! tip "Key Takeaways"
|
||||
|
||||
- RF-DETR-2XL achieves 60.1 AP50:95 on COCO detection at 17.2 ms latency (T4, TensorRT FP16)
|
||||
- RF-DETR-L outperforms YOLOv11x (56.5 vs 50.9 AP50:95) at lower latency (6.8 vs 10.5 ms)
|
||||
- Segmentation models range from 40.3 AP (Nano, 3.4 ms) to 49.9 AP (2XL, 21.8 ms) on COCO
|
||||
- All latency measured on NVIDIA T4 with TensorRT 10.4, CUDA 12.4, FP16, batch size 1
|
||||
- RF100-VL results demonstrate strong domain-shift generalization across 100 diverse datasets
|
||||
- RF-DETR Keypoint (Preview) achieves 71.8 AP50:95 on COCO person keypoints at 9.7 ms (T4, TensorRT FP16)
|
||||
|
||||
This page reports RF-DETR benchmark results for object detection, instance segmentation, and keypoint detection on Microsoft COCO and RF100-VL (detection only). All benchmark numbers and plots match the latest released checkpoints and tables shown below. Latency values are measured on an NVIDIA T4 with TensorRT in FP16 at batch size 1. For full methodology details and architectural context, see the RF-DETR paper.
|
||||
|
||||
## Methodology
|
||||
|
||||
Accuracy is reported using standard COCO metrics computed with pycocotools. For object detection, we report COCO AP50 and COCO AP50:95, and the same metrics are also reported for RF100-VL. COCO results are evaluated on the validation split, following common practice in detector benchmarking. RF100-VL results are averaged across all 100 datasets to reflect performance under diverse real-world data distributions.
|
||||
|
||||
Latency is measured as single-image inference latency rather than sustained throughput. All latency numbers are obtained on an NVIDIA T4 GPU using TensorRT 10.4 and CUDA 12.4 with FP16 inference and batch size 1. To reduce variance caused by GPU power throttling and thermal effects, a 200 ms buffer is inserted between consecutive forward passes. This procedure improves reproducibility of latency measurements but is not intended to measure maximum throughput.
|
||||
|
||||
Accuracy and latency are always measured using the same model artifact and the same numerical precision. This avoids reporting FP32 accuracy together with FP16 latency, which can lead to misleading comparisons because naive FP16 conversion can significantly degrade accuracy for some models.
|
||||
|
||||
!!! info "Metric definitions"
|
||||
|
||||
**AP50**: Detection accuracy at IoU threshold >= 0.50.
|
||||
**AP50:95**: Mean accuracy averaged over IoU thresholds 0.50 to 0.95 (step 0.05) — the primary COCO metric.
|
||||
Latency measured on NVIDIA T4, TensorRT 10.4, CUDA 12.4, FP16, batch size 1,
|
||||
with 200 ms thermal buffer between passes to reduce GPU thermal variance.
|
||||
|
||||
## Detection
|
||||
|
||||
<img alt="rf_detr_1-4_latency_accuracy_object_detection" src="https://storage.googleapis.com/com-roboflow-marketing/rf-detr/rf_detr_1-4_latency_accuracy_object_detection.png" />
|
||||
|
||||
| Architecture | COCO AP<sub>50</sub> | COCO AP<sub>50:95</sub> | RF100VL AP<sub>50</sub> | RF100VL AP<sub>50:95</sub> | Latency (ms) | Params (M) | Resolution |
|
||||
| :----------: | :------------------: | :---------------------: | :---------------------: | :------------------------: | :----------: | :--------: | :--------: |
|
||||
| RF-DETR-N | 67.6 | 48.4 | 85.0 | 57.7 | 2.3 | 30.5 | 384x384 |
|
||||
| RF-DETR-S | 72.1 | 53.0 | 86.7 | 60.2 | 3.5 | 32.1 | 512x512 |
|
||||
| RF-DETR-M | 73.6 | 54.7 | 87.4 | 61.2 | 4.4 | 33.7 | 576x576 |
|
||||
| RF-DETR-L | 75.1 | 56.5 | 88.2 | 62.2 | 6.8 | 33.9 | 704x704 |
|
||||
| RF-DETR-XL | 77.4 | 58.6 | 88.5 | 62.9 | 11.5 | 126.4 | 700x700 |
|
||||
| RF-DETR-2XL | 78.5 | 60.1 | 89.0 | 63.2 | 17.2 | 126.9 | 880x880 |
|
||||
| YOLO11-N | 52.0 | 37.4 | 81.4 | 55.3 | 2.5 | 2.6 | 640x640 |
|
||||
| YOLO11-S | 59.7 | 44.4 | 82.3 | 56.2 | 3.2 | 9.4 | 640x640 |
|
||||
| YOLO11-M | 64.1 | 48.6 | 82.5 | 56.5 | 5.1 | 20.1 | 640x640 |
|
||||
| YOLO11-L | 64.9 | 49.9 | 82.2 | 56.5 | 6.5 | 25.3 | 640x640 |
|
||||
| YOLO11-X | 66.1 | 50.9 | 81.7 | 56.2 | 10.5 | 56.9 | 640x640 |
|
||||
| YOLO26-N | 55.8 | 40.3 | 76.7 | 52.0 | 1.7 | 2.6 | 640x640 |
|
||||
| YOLO26-S | 64.3 | 47.7 | 82.7 | 57.0 | 2.6 | 9.4 | 640x640 |
|
||||
| YOLO26-M | 69.7 | 52.5 | 84.4 | 58.7 | 4.4 | 20.1 | 640x640 |
|
||||
| YOLO26-L | 71.1 | 54.1 | 85.0 | 59.3 | 5.7 | 25.3 | 640x640 |
|
||||
| YOLO26-X | 74.0 | 56.9 | 85.6 | 60.0 | 9.6 | 56.9 | 640x640 |
|
||||
| LW-DETR-T | 60.7 | 42.9 | 84.7 | 57.1 | 1.9 | 12.1 | 640x640 |
|
||||
| LW-DETR-S | 66.8 | 48.0 | 85.0 | 57.4 | 2.6 | 14.6 | 640x640 |
|
||||
| LW-DETR-M | 72.0 | 52.6 | 86.8 | 59.8 | 4.4 | 28.2 | 640x640 |
|
||||
| LW-DETR-L | 74.6 | 56.1 | 87.4 | 61.5 | 6.9 | 46.8 | 640x640 |
|
||||
| LW-DETR-X | 76.9 | 58.3 | 87.9 | 62.1 | 13.0 | 118.0 | 640x640 |
|
||||
| D-FINE-N | 60.2 | 42.7 | 84.4 | 58.2 | 2.1 | 3.8 | 640x640 |
|
||||
| D-FINE-S | 67.6 | 50.6 | 85.3 | 60.3 | 3.5 | 10.2 | 640x640 |
|
||||
| D-FINE-M | 72.6 | 55.0 | 85.5 | 60.6 | 5.4 | 19.2 | 640x640 |
|
||||
| D-FINE-L | 74.9 | 57.2 | 86.4 | 61.6 | 7.5 | 31.0 | 640x640 |
|
||||
| D-FINE-X | 76.8 | 59.3 | 86.9 | 62.2 | 11.5 | 62.0 | 640x640 |
|
||||
|
||||
## Segmentation
|
||||
|
||||
<img alt="rf_detr_1-4_latency_accuracy_instance_segmentation" src="https://storage.googleapis.com/com-roboflow-marketing/rf-detr/rf_detr_1-4_latency_accuracy_instance_segmentation.png" />
|
||||
|
||||
| Architecture | COCO AP<sub>50</sub> | COCO AP<sub>50:95</sub> | Latency (ms) | Params (M) | Resolution |
|
||||
| :-------------: | :------------------: | :---------------------: | :----------: | :--------: | :--------: |
|
||||
| RF-DETR-Seg-N | 63.0 | 40.3 | 3.4 | 33.6 | 312x312 |
|
||||
| RF-DETR-Seg-S | 66.2 | 43.1 | 4.4 | 33.7 | 384x384 |
|
||||
| RF-DETR-Seg-M | 68.4 | 45.3 | 5.9 | 35.7 | 432x432 |
|
||||
| RF-DETR-Seg-L | 70.5 | 47.1 | 8.8 | 36.2 | 504x504 |
|
||||
| RF-DETR-Seg-XL | 72.2 | 48.8 | 13.5 | 38.1 | 624x624 |
|
||||
| RF-DETR-Seg-2XL | 73.1 | 49.9 | 21.8 | 38.6 | 768x768 |
|
||||
| YOLOv8-N-Seg | 45.6 | 28.3 | 3.5 | 3.4 | 640x640 |
|
||||
| YOLOv8-S-Seg | 53.8 | 34.0 | 4.2 | 11.8 | 640x640 |
|
||||
| YOLOv8-M-Seg | 58.2 | 37.3 | 7.0 | 27.3 | 640x640 |
|
||||
| YOLOv8-L-Seg | 60.5 | 39.0 | 9.7 | 46.0 | 640x640 |
|
||||
| YOLOv8-XL-Seg | 61.3 | 39.5 | 14.0 | 71.8 | 640x640 |
|
||||
| YOLOv11-N-Seg | 47.8 | 30.0 | 3.6 | 2.9 | 640x640 |
|
||||
| YOLOv11-S-Seg | 55.4 | 35.0 | 4.6 | 10.1 | 640x640 |
|
||||
| YOLOv11-M-Seg | 60.0 | 38.5 | 6.9 | 22.4 | 640x640 |
|
||||
| YOLOv11-L-Seg | 61.5 | 39.5 | 8.3 | 27.6 | 640x640 |
|
||||
| YOLOv11-XL-Seg | 62.4 | 40.1 | 13.7 | 62.1 | 640x640 |
|
||||
| YOLO26-N-Seg | 54.3 | 34.7 | 2.31 | 2.7 | 640x640 |
|
||||
| YOLO26-S-Seg | 62.4 | 40.2 | 3.47 | 10.4 | 640x640 |
|
||||
| YOLO26-M-Seg | 67.8 | 44.0 | 6.32 | 23.6 | 640x640 |
|
||||
| YOLO26-L-Seg | 69.8 | 45.5 | 7.58 | 28.0 | 640x640 |
|
||||
| YOLO26-X-Seg | 71.6 | 46.8 | 12.92 | 62.8 | 640x640 |
|
||||
|
||||
## Keypoints
|
||||
|
||||
<img alt="RF-DETR Keypoint mAP vs latency chart comparing against YOLO26-pose and YOLO11-pose on MS COCO" src="../assets/keypoints/kp-map-latency.png" />
|
||||
|
||||
| Architecture | COCO AP<sub>50:95</sub> | Latency (ms) |
|
||||
| :------------------------: | :---------------------: | :----------: |
|
||||
| RF-DETR Keypoint (Preview) | 71.8 | 9.7 |
|
||||
| YOLO11-pose N | 48.9 | 3.2 |
|
||||
| YOLO11-pose S | 57.5 | 3.4 |
|
||||
| YOLO11-pose M | 64.2 | 5.2 |
|
||||
| YOLO11-pose L | 65.2 | 6.6 |
|
||||
| YOLO11-pose X | 68.6 | 10.6 |
|
||||
| YOLO26-pose N | 55.9 | 1.9 |
|
||||
| YOLO26-pose S | 62.0 | 2.7 |
|
||||
| YOLO26-pose M | 68.0 | 4.6 |
|
||||
| YOLO26-pose L | 69.2 | 5.9 |
|
||||
| YOLO26-pose X | 71.0 | 9.8 |
|
||||
|
||||
> Keypoint benchmarks report AP<sub>50:95</sub> (OKS-based); this is the standard COCO keypoint comparison metric.
|
||||
@@ -0,0 +1,115 @@
|
||||
---
|
||||
description: Deploy fine-tuned RF-DETR detection and segmentation models to Roboflow for cloud inference, edge hardware, and multi-step vision workflows.
|
||||
---
|
||||
|
||||
# Deploy a Trained RF-DETR Model
|
||||
|
||||
!!! tip "Key Takeaways"
|
||||
|
||||
- Deploy fine-tuned RF-DETR models to Roboflow with a single `deploy_to_roboflow()` call
|
||||
- Supports both detection and segmentation model deployment
|
||||
- Run deployed models via Roboflow Inference on cloud, edge hardware, or NVIDIA Jetson
|
||||
- Model weights are cached locally after the first inference run for fast subsequent predictions
|
||||
|
||||
You can deploy a fine-tuned RF-DETR model to Roboflow.
|
||||
|
||||
Deploying to Roboflow allows you to create multi-step computer vision applications that run both in the cloud and your own hardware.
|
||||
|
||||
To deploy your model to Roboflow, run:
|
||||
|
||||
=== "Object Detection"
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRNano
|
||||
|
||||
x = RFDETRNano(pretrain_weights="<path/to/pretrain/weights/dir>")
|
||||
x.deploy_to_roboflow(
|
||||
workspace="<your-workspace>",
|
||||
project_id="<your-project-id>",
|
||||
version=1,
|
||||
api_key="<YOUR_API_KEY>",
|
||||
)
|
||||
```
|
||||
|
||||
=== "Image Segmentation"
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRSegMedium
|
||||
|
||||
x = RFDETRSegMedium(pretrain_weights="<path/to/pretrain/weights/dir>")
|
||||
x.deploy_to_roboflow(
|
||||
workspace="<your-workspace>",
|
||||
project_id="<your-project-id>",
|
||||
version=1,
|
||||
api_key="<YOUR_API_KEY>",
|
||||
)
|
||||
```
|
||||
|
||||
Above, set your Roboflow Workspace ID, the ID of the project to which you want to upload your model, and your Roboflow API key.
|
||||
|
||||
- [Learn how to find your Workspace and Project ID.](https://docs.roboflow.com/developer/authentication/workspace-and-project-ids)
|
||||
- [Learn how to find your API key.](https://docs.roboflow.com/developer/authentication/find-your-roboflow-api-key)
|
||||
|
||||
You can then run your model with Roboflow Inference:
|
||||
|
||||
=== "Object Detection"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import requests
|
||||
|
||||
url = "https://media.roboflow.com/dog.jpeg"
|
||||
image = Image.open(BytesIO(requests.get(url).content))
|
||||
|
||||
model = get_model("rfdetr-large") # replace with your Roboflow model ID
|
||||
|
||||
predictions = model.infer(image, confidence=0.5)[0]
|
||||
|
||||
detections = sv.Detections.from_inference(predictions)
|
||||
|
||||
labels = [prediction.class_name for prediction in predictions.predictions]
|
||||
|
||||
annotated_image = image.copy()
|
||||
annotated_image = sv.BoxAnnotator().annotate(annotated_image, detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
|
||||
|
||||
sv.plot_image(annotated_image)
|
||||
```
|
||||
|
||||
=== "Image Segmentation"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import requests
|
||||
|
||||
url = "https://media.roboflow.com/dog.jpeg"
|
||||
image = Image.open(BytesIO(requests.get(url).content))
|
||||
|
||||
model = get_model("rfdetr-seg-small") # replace with your Roboflow model ID
|
||||
|
||||
predictions = model.infer(image, confidence=0.5)[0]
|
||||
|
||||
detections = sv.Detections.from_inference(predictions)
|
||||
|
||||
labels = [prediction.class_name for prediction in predictions.predictions]
|
||||
|
||||
annotated_image = image.copy()
|
||||
annotated_image = sv.MaskAnnotator().annotate(annotated_image, detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
|
||||
|
||||
sv.plot_image(annotated_image)
|
||||
```
|
||||
|
||||
Above, replace `rfdetr-large` with the your Roboflow model ID. You can find this ID from the "Models" list in your Roboflow dashboard:
|
||||
|
||||

|
||||
|
||||
When you first run this model, your model weights will be cached for local use with Inference.
|
||||
|
||||
You will then see the results from your fine-tuned model.
|
||||
@@ -0,0 +1,396 @@
|
||||
---
|
||||
description: Export RF-DETR models to ONNX, TensorRT, and TFLite (FP32/FP16/INT8) for high-performance inference on GPUs, mobile, and edge devices.
|
||||
---
|
||||
|
||||
# Export RF-DETR Model
|
||||
|
||||
!!! tip "Key Takeaways"
|
||||
|
||||
- Export to ONNX for cross-platform inference with ONNX Runtime, OpenVINO, or TensorRT
|
||||
- Export to TFLite (FP32, FP16, INT8) for mobile and edge deployment
|
||||
- TensorRT conversion delivers lowest latency on NVIDIA GPUs (2.3 ms for Nano)
|
||||
- INT8 quantization requires calibration data from your dataset for accurate results
|
||||
- Custom input resolutions supported (must be divisible by `patch_size × num_windows`, which varies by model variant)
|
||||
|
||||
RF-DETR supports exporting models to ONNX and TFLite formats, enabling deployment across a wide range of inference frameworks, edge devices, and hardware accelerators.
|
||||
|
||||
## Installation
|
||||
|
||||
Install the export dependencies you need:
|
||||
|
||||
```bash
|
||||
# ONNX export only
|
||||
pip install "rfdetr[onnx]"
|
||||
|
||||
# TFLite export (includes ONNX dependency)
|
||||
pip install "rfdetr[onnx,tflite]"
|
||||
```
|
||||
|
||||
## Basic Export
|
||||
|
||||
Export your trained model to ONNX format:
|
||||
|
||||
=== "Object Detection"
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRMedium
|
||||
|
||||
model = RFDETRMedium(pretrain_weights="<path/to/checkpoint.pth>")
|
||||
|
||||
model.export()
|
||||
```
|
||||
|
||||
=== "Image Segmentation"
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRSegMedium
|
||||
|
||||
model = RFDETRSegMedium(pretrain_weights="<path/to/checkpoint.pth>")
|
||||
|
||||
model.export()
|
||||
```
|
||||
|
||||
This command saves the ONNX model to the `output` directory by default.
|
||||
|
||||
## Export Parameters
|
||||
|
||||
The `export()` method accepts several parameters to customize the export process:
|
||||
|
||||
| Parameter | Default | Description |
|
||||
| ------------------ | ---------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `output_dir` | `"output"` | Directory where the exported model will be saved. |
|
||||
| `format` | `"onnx"` | Export format: `"onnx"` or `"tflite"`. |
|
||||
| `quantization` | `None` | TFLite quantization mode: `None`/`"fp32"`, `"fp16"`, or `"int8"`. Only used when `format="tflite"`. |
|
||||
| `calibration_data` | `None` | Calibration data for TFLite export. Image directory, `.npy` file path, NumPy array, or `None`. See [TFLite Export](#tflite-export). |
|
||||
| `max_images` | `100` | Maximum number of images to load from a calibration directory for TFLite INT8 quantization. Ignored for other calibration data formats. |
|
||||
| `infer_dir` | `None` | Optional directory of sample images for inference validation during export tracing. If not provided, a random dummy image is generated. |
|
||||
| `backbone_only` | `False` | Export only the backbone feature extractor instead of the full model. |
|
||||
| `opset_version` | `17` | ONNX opset version to use for export. Higher versions support more operations. |
|
||||
| `verbose` | `True` | Whether to print verbose export information. |
|
||||
| `shape` | `None` | Input shape as tuple `(height, width)`. Each dimension must be divisible by the selected model's block size (`patch_size * num_windows`). If not provided, uses the model's default resolution. |
|
||||
| `batch_size` | `1` | Batch size for the exported model. |
|
||||
| `dynamic_batch` | `False` | If `True`, export with a dynamic batch dimension so the ONNX model accepts variable batch sizes at runtime. |
|
||||
| `patch_size` | `None` | Backbone patch size override. Defaults to the value from `model_config.patch_size`. Must match the instantiated model's patch size when provided. |
|
||||
| `notes` | `None` | Optional user-defined metadata (string, dict, list, or any JSON-serialisable value) to embed in the exported ONNX model under the `"rfdetr_notes"` metadata property. |
|
||||
|
||||
## Advanced Export Examples
|
||||
|
||||
### Export with Custom Output Directory
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRMedium
|
||||
|
||||
model = RFDETRMedium(pretrain_weights="<path/to/checkpoint.pth>")
|
||||
|
||||
model.export(output_dir="exports/my_model")
|
||||
```
|
||||
|
||||
### Export with Custom Resolution
|
||||
|
||||
Export the model with a specific input resolution. For example, `RFDETRMedium` expects dimensions divisible by `32` (`patch_size=16`, `num_windows=2`):
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRMedium
|
||||
|
||||
model = RFDETRMedium(pretrain_weights="<path/to/checkpoint.pth>")
|
||||
|
||||
model.export(shape=(608, 608))
|
||||
```
|
||||
|
||||
### Export Backbone Only
|
||||
|
||||
Export only the backbone feature extractor for use in custom pipelines:
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRMedium
|
||||
|
||||
model = RFDETRMedium(pretrain_weights="<path/to/checkpoint.pth>")
|
||||
|
||||
model.export(backbone_only=True)
|
||||
```
|
||||
|
||||
## Output Files
|
||||
|
||||
After running the export, you will find the following files in your output directory:
|
||||
|
||||
- `inference_model.onnx` - The exported ONNX model (or `backbone_model.onnx` if `backbone_only=True`)
|
||||
|
||||
## Optional: Convert ONNX to TensorRT
|
||||
|
||||
If you want lower latency on NVIDIA GPUs, you can convert the exported ONNX model to a TensorRT engine.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> Run TensorRT conversion on the same machine and GPU family where you plan to deploy inference.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Install TensorRT (`trtexec` must be available in your `PATH`)
|
||||
- Export an ONNX model first (for example: `output/inference_model.onnx`)
|
||||
|
||||
### Python API Conversion
|
||||
|
||||
```python
|
||||
from argparse import Namespace
|
||||
|
||||
from rfdetr.export._tensorrt import trtexec
|
||||
|
||||
args = Namespace(
|
||||
verbose=True,
|
||||
profile=False,
|
||||
dry_run=False,
|
||||
)
|
||||
|
||||
trtexec("output/inference_model.onnx", args)
|
||||
```
|
||||
|
||||
This produces `output/inference_model.engine`. If `profile=True`, it also writes an Nsight Systems report (`.nsys-rep`).
|
||||
|
||||
## TFLite Export
|
||||
|
||||
!!! warning "Experimental — Use with Caution"
|
||||
|
||||
TFLite export is **experimental and work-in-progress**. The pipeline depends on
|
||||
several upstream packages (`onnx2tf`, `ai_edge_litert`, `tflite-runtime`) that
|
||||
have experienced breaking API changes and installation instabilities across
|
||||
releases. You may encounter errors or unexpected results.
|
||||
|
||||
**Known instabilities:**
|
||||
|
||||
- `onnx2tf` output graph structure can change between minor versions, silently
|
||||
altering output tensor layout and breaking downstream inference code.
|
||||
- `ai_edge_litert` (Google's replacement for `tflite-runtime`) is still
|
||||
stabilising its public API; version pinning is strongly recommended.
|
||||
- INT8 quantization accuracy is sensitive to calibration data quality — poor
|
||||
calibration causes silent precision loss with no error at export time.
|
||||
- The ONNX → TF → TFLite conversion chain introduces numerical rounding that
|
||||
may produce slightly different predictions from the original PyTorch model.
|
||||
- Installation of the `[tflite]` extra may conflict with existing TensorFlow
|
||||
or NumPy versions in your environment.
|
||||
|
||||
**Recommendations:**
|
||||
|
||||
- Pin your dependency versions (e.g. `onnx2tf==X.Y.Z`) and test before each upgrade.
|
||||
- Validate exported `.tflite` files against a held-out evaluation set before deploying.
|
||||
- Prefer ONNX export when your target runtime supports it — it is more stable and
|
||||
better tested.
|
||||
- If export fails, check the [open issues](https://github.com/roboflow/rf-detr/issues)
|
||||
for known workarounds or report a new one with your environment details
|
||||
(`pip freeze`, Python version, OS).
|
||||
|
||||
Export your model to TFLite for deployment on mobile devices, microcontrollers, and edge hardware via TensorFlow Lite. The TFLite export pipeline converts ONNX → TensorFlow → TFLite using [onnx2tf](https://github.com/PINTO0309/onnx2tf).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
pip install "rfdetr[onnx,tflite]"
|
||||
```
|
||||
|
||||
### Basic TFLite Export (FP32)
|
||||
|
||||
=== "Object Detection"
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRSmall
|
||||
|
||||
model = RFDETRSmall()
|
||||
|
||||
model.export(format="tflite", output_dir="output")
|
||||
```
|
||||
|
||||
=== "Image Segmentation"
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRSegNano
|
||||
|
||||
model = RFDETRSegNano()
|
||||
|
||||
model.export(format="tflite", output_dir="output")
|
||||
```
|
||||
|
||||
This produces both `output/inference_model_float32.tflite` and `output/inference_model_float16.tflite`.
|
||||
|
||||
### INT8 Quantization with Calibration Data
|
||||
|
||||
For INT8 quantization, provide representative images from your dataset as calibration data. This is **critical** for preserving model accuracy — without real calibration data, the quantizer uses random noise and accuracy will be poor.
|
||||
|
||||
#### Option 1: Point to an Image Directory (Recommended)
|
||||
|
||||
The simplest approach — just point `calibration_data` to a directory containing JPEG/PNG images. The converter automatically loads, resizes, and prepares the images:
|
||||
|
||||
```python
|
||||
from rfdetr import RFDETRNano
|
||||
|
||||
model = RFDETRNano()
|
||||
model.export(
|
||||
format="tflite",
|
||||
quantization="int8",
|
||||
calibration_data="path/to/val2017/", # directory of images
|
||||
output_dir="output",
|
||||
)
|
||||
```
|
||||
|
||||
The converter loads up to 100 images from the directory by default, resizes them to the model's input resolution, and uses them for both output validation and INT8 calibration. Supported formats: JPEG, PNG, BMP, WebP.
|
||||
|
||||
You can control how many images are loaded with the `max_images` parameter:
|
||||
|
||||
```python
|
||||
model.export(
|
||||
format="tflite",
|
||||
quantization="int8",
|
||||
calibration_data="path/to/val2017/",
|
||||
max_images=200, # load up to 200 images (default: 100)
|
||||
output_dir="output",
|
||||
)
|
||||
```
|
||||
|
||||
#### Option 2: NumPy `.npy` File
|
||||
|
||||
Prepare calibration data as a NumPy array and save it to a `.npy` file:
|
||||
|
||||
- Shape: `(N, H, W, 3)` — NHWC format with 3 color channels
|
||||
- Data type: `float32`
|
||||
- Value range: `[0, 1]` (divide by 255, but do **not** apply ImageNet normalization — the converter handles that automatically)
|
||||
- Recommended: 20–100 representative images from your dataset
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from rfdetr import RFDETRSmall
|
||||
|
||||
model = RFDETRSmall()
|
||||
target_resolution = model.model_config.resolution
|
||||
|
||||
# Load representative images from your dataset
|
||||
images = []
|
||||
for path in image_paths[:50]: # 50 representative samples
|
||||
img = Image.open(path).convert("RGB").resize((target_resolution, target_resolution))
|
||||
images.append(np.array(img, dtype=np.float32) / 255.0)
|
||||
|
||||
calibration_data = np.stack(images) # shape: (50, H, W, 3)
|
||||
|
||||
# Save to .npy for reuse
|
||||
np.save("calibration_data.npy", calibration_data)
|
||||
|
||||
# Export with INT8 quantization
|
||||
model.export(
|
||||
format="tflite",
|
||||
quantization="int8",
|
||||
calibration_data="calibration_data.npy",
|
||||
output_dir="output",
|
||||
)
|
||||
```
|
||||
|
||||
#### Option 3: NumPy Array Directly
|
||||
|
||||
You can also pass the NumPy array directly without saving to disk:
|
||||
|
||||
```python
|
||||
model.export(
|
||||
format="tflite",
|
||||
quantization="int8",
|
||||
calibration_data=calibration_data, # np.ndarray
|
||||
output_dir="output",
|
||||
)
|
||||
```
|
||||
|
||||
### FP16 Export
|
||||
|
||||
FP16 models are always produced alongside FP32. You can explicitly request FP16 mode:
|
||||
|
||||
```python
|
||||
model.export(format="tflite", quantization="fp16", output_dir="output")
|
||||
```
|
||||
|
||||
### TFLite Output Files
|
||||
|
||||
The `onnx2tf` converter **always** produces both FP32 and FP16 TFLite files, regardless of the requested quantization mode. When `quantization="int8"` is specified, it additionally produces the INT8-quantized model.
|
||||
|
||||
| File | Description |
|
||||
| -------------------------------------- | --------------------------------------- |
|
||||
| `inference_model_float32.tflite` | FP32 model (always produced) |
|
||||
| `inference_model_float16.tflite` | FP16 model (always produced) |
|
||||
| `inference_model_integer_quant.tflite` | INT8 model (when `quantization="int8"`) |
|
||||
|
||||
!!! note
|
||||
|
||||
Segmentation models produce TFLite files with three outputs: `dets` (bounding boxes), `labels` (class scores), and `masks` (per-instance segmentation masks).
|
||||
|
||||
### TFLite Inference Example
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# pip install tflite-runtime (or use tensorflow.lite)
|
||||
import tflite_runtime.interpreter as tflite
|
||||
|
||||
# Load model
|
||||
interpreter = tflite.Interpreter(model_path="output/inference_model_float32.tflite")
|
||||
interpreter.allocate_tensors()
|
||||
|
||||
input_details = interpreter.get_input_details()
|
||||
output_details = interpreter.get_output_details()
|
||||
|
||||
# Prepare input — TFLite model expects NHWC, ImageNet-normalized
|
||||
input_height, input_width = input_details[0]["shape"][1:3]
|
||||
image = Image.open("image.jpg").convert("RGB").resize((input_width, input_height))
|
||||
image_array = np.array(image, dtype=np.float32) / 255.0
|
||||
|
||||
# Apply ImageNet normalization
|
||||
mean = np.array([0.485, 0.456, 0.406])
|
||||
std = np.array([0.229, 0.224, 0.225])
|
||||
image_array = (image_array - mean) / std
|
||||
|
||||
# Add batch dimension: (1, H, W, 3)
|
||||
image_array = np.expand_dims(image_array, axis=0).astype(np.float32)
|
||||
|
||||
# Run inference
|
||||
interpreter.set_tensor(input_details[0]["index"], image_array)
|
||||
interpreter.invoke()
|
||||
|
||||
boxes = interpreter.get_tensor(output_details[0]["index"])
|
||||
labels = interpreter.get_tensor(output_details[1]["index"])
|
||||
```
|
||||
|
||||
## Using the Exported Model
|
||||
|
||||
Once exported, you can use the ONNX model with various inference frameworks:
|
||||
|
||||
### ONNX Runtime
|
||||
|
||||
```python
|
||||
import onnxruntime as ort
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# Load the ONNX model
|
||||
session = ort.InferenceSession("output/inference_model.onnx")
|
||||
|
||||
# Prepare input image
|
||||
input_height, input_width = session.get_inputs()[0].shape[2:4]
|
||||
image = Image.open("image.jpg").convert("RGB")
|
||||
image = image.resize((input_width, input_height)) # Resize to the exported model shape
|
||||
image_array = np.array(image).astype(np.float32) / 255.0
|
||||
|
||||
# Normalize
|
||||
mean = np.array([0.485, 0.456, 0.406])
|
||||
std = np.array([0.229, 0.224, 0.225])
|
||||
image_array = (image_array - mean) / std
|
||||
|
||||
# Convert to NCHW format
|
||||
image_array = np.transpose(image_array, (2, 0, 1))
|
||||
image_array = np.expand_dims(image_array, axis=0)
|
||||
|
||||
# Run inference
|
||||
outputs = session.run(None, {"input": image_array})
|
||||
boxes, labels = outputs
|
||||
```
|
||||
|
||||
## Next Steps
|
||||
|
||||
After exporting your model, you may want to:
|
||||
|
||||
- [Deploy to Roboflow](deploy.md) for cloud-based inference and workflow integration
|
||||
- Use the ONNX model with TensorRT for optimized GPU inference
|
||||
- Deploy TFLite models on mobile/edge devices with TensorFlow Lite
|
||||
- Integrate with edge deployment frameworks like ONNX Runtime or OpenVINO
|
||||
@@ -0,0 +1,198 @@
|
||||
---
|
||||
description: Run pre-trained RF-DETR models (Nano to 2XLarge) on images, video, webcam, and RTSP streams. COCO-trained with real-time DINOv2 backbone.
|
||||
---
|
||||
|
||||
You can run any of the four supported RF-DETR base models -- Nano, Small, Medium, Large -- with [Inference](https://github.com/roboflow/inference), an open source computer vision inference server. The base models are trained on the [Microsoft COCO dataset](https://universe.roboflow.com/microsoft/coco). XLarge and 2XLarge detection models are also available via `pip install rfdetr[plus]` and are provided under the PML 1.0 license.
|
||||
|
||||
=== "Run on an Image"
|
||||
|
||||
To run RF-DETR on an image, use the following code:
|
||||
|
||||
```python
|
||||
import os
|
||||
import supervision as sv
|
||||
from inference import get_model
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import requests
|
||||
|
||||
url = "https://media.roboflow.com/dog.jpeg"
|
||||
image = Image.open(BytesIO(requests.get(url).content))
|
||||
|
||||
model = get_model("rfdetr-large")
|
||||
|
||||
predictions = model.infer(image, confidence=0.5)[0]
|
||||
|
||||
detections = sv.Detections.from_inference(predictions)
|
||||
|
||||
labels = [prediction.class_name for prediction in predictions.predictions]
|
||||
|
||||
annotated_image = image.copy()
|
||||
annotated_image = sv.BoxAnnotator().annotate(annotated_image, detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
|
||||
|
||||
sv.plot_image(annotated_image)
|
||||
```
|
||||
|
||||
Above, replace the image URL with any image you want to use with the model.
|
||||
|
||||
Here are the results from the code above:
|
||||
|
||||
<figure markdown="span">
|
||||
{ width=300 }
|
||||
<figcaption>RF-DETR Base predictions</figcaption>
|
||||
</figure>
|
||||
|
||||
=== "Run on a Video File"
|
||||
|
||||
To run RF-DETR on a video file, use the following code:
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
|
||||
def callback(frame, index):
|
||||
detections = model.predict(frame[:, :, ::-1], threshold=0.5)
|
||||
|
||||
labels = [
|
||||
f"{COCO_CLASSES[class_id]} {confidence:.2f}"
|
||||
for class_id, confidence in zip(detections.class_id, detections.confidence)
|
||||
]
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = sv.BoxAnnotator().annotate(annotated_frame, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
return annotated_frame
|
||||
|
||||
|
||||
sv.process_video(
|
||||
source_path="<SOURCE_VIDEO_PATH>",
|
||||
target_path="<TARGET_VIDEO_PATH>",
|
||||
callback=callback,
|
||||
)
|
||||
```
|
||||
|
||||
Above, set your `SOURCE_VIDEO_PATH` and `TARGET_VIDEO_PATH` to the directories of the video you want to process and where you want to save the results from inference, respectively.
|
||||
|
||||
=== "Run on a Webcam Stream"
|
||||
|
||||
To run RF-DETR on a webcam input, use the following code:
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
cap = cv2.VideoCapture(0)
|
||||
while True:
|
||||
success, frame = cap.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
detections = model.predict(frame[:, :, ::-1], threshold=0.5)
|
||||
|
||||
labels = [
|
||||
f"{COCO_CLASSES[class_id]} {confidence:.2f}"
|
||||
for class_id, confidence in zip(detections.class_id, detections.confidence)
|
||||
]
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = sv.BoxAnnotator().annotate(annotated_frame, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("Webcam", annotated_frame)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "Run on an RTSP Stream"
|
||||
|
||||
To run RF-DETR on an RTSP (Real Time Streaming Protocol) stream, use the following code:
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
cap = cv2.VideoCapture("<RTSP_STREAM_URL>")
|
||||
while True:
|
||||
success, frame = cap.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
detections = model.predict(frame[:, :, ::-1], threshold=0.5)
|
||||
|
||||
labels = [
|
||||
f"{COCO_CLASSES[class_id]} {confidence:.2f}"
|
||||
for class_id, confidence in zip(detections.class_id, detections.confidence)
|
||||
]
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = sv.BoxAnnotator().annotate(annotated_frame, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RTSP Stream", annotated_frame)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
You can change the RF-DETR model that the code snippet above uses. To do so, update `rfdetr-base` to any of the following values:
|
||||
|
||||
- `rfdetr-nano`
|
||||
- `rfdetr-small`
|
||||
- `rfdetr-medium`
|
||||
- `rfdetr-large`
|
||||
|
||||
## Batch Inference
|
||||
|
||||
You can provide `.predict()` with either a single image or a list of images. When multiple images are supplied, they are processed together in a single forward pass, resulting in a corresponding list of detections.
|
||||
|
||||
```python
|
||||
import io
|
||||
import requests
|
||||
import supervision as sv
|
||||
from PIL import Image
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
urls = [
|
||||
"https://media.roboflow.com/notebooks/examples/dog-2.jpeg",
|
||||
"https://media.roboflow.com/notebooks/examples/dog-3.jpeg",
|
||||
]
|
||||
|
||||
images = [Image.open(io.BytesIO(requests.get(url).content)) for url in urls]
|
||||
|
||||
detections_list = model.predict(images, threshold=0.5)
|
||||
|
||||
for image, detections in zip(images, detections_list):
|
||||
labels = [
|
||||
f"{COCO_CLASSES[class_id]} {confidence:.2f}"
|
||||
for class_id, confidence in zip(detections.class_id, detections.confidence)
|
||||
]
|
||||
|
||||
annotated_image = image.copy()
|
||||
annotated_image = sv.BoxAnnotator().annotate(annotated_image, detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
|
||||
|
||||
sv.plot_image(annotated_image)
|
||||
```
|
||||
@@ -0,0 +1,185 @@
|
||||
---
|
||||
description: Run RF-DETR object detection on images, video, and streams. Nano to 2XLarge models with 2.3-17.2 ms latency and up to 60.1 AP on COCO.
|
||||
---
|
||||
|
||||
# Run an RF-DETR Object Detection Model
|
||||
|
||||
RF-DETR is a real-time transformer architecture for object detection, built on a DINOv2 vision transformer backbone. The base models are trained on the Microsoft COCO dataset and achieve state-of-the-art accuracy and latency trade-offs.
|
||||
|
||||
## Pre-trained Checkpoints
|
||||
|
||||
RF-DETR offers model sizes from Nano to 2XLarge, allowing trade-offs between accuracy, latency, and parameter count. All latency numbers were measured on an NVIDIA T4 using TensorRT, FP16, and batch size 1. Core models (Nano to Large) are licensed under Apache 2.0. XLarge and 2XLarge (marked with △) are provided by the [`rfdetr_plus`](https://github.com/roboflow/rf-detr-plus) extension (`pip install rfdetr[plus]`) under the Platform Model License 1.0 and require a Roboflow account.
|
||||
|
||||
| Size | RF-DETR package class | Inference package alias | COCO AP<sub>50</sub> | COCO AP<sub>50:95</sub> | Latency (ms) | Params (M) | Resolution | License |
|
||||
| :--: | :-------------------: | :---------------------- | :------------------: | :---------------------: | :----------: | :--------: | :--------: | :--------: |
|
||||
| N | `RFDETRNano` | `rfdetr-nano` | 67.6 | 48.4 | 2.3 | 30.5 | 384x384 | Apache 2.0 |
|
||||
| S | `RFDETRSmall` | `rfdetr-small` | 72.1 | 53.0 | 3.5 | 32.1 | 512x512 | Apache 2.0 |
|
||||
| M | `RFDETRMedium` | `rfdetr-medium` | 73.6 | 54.7 | 4.4 | 33.7 | 576x576 | Apache 2.0 |
|
||||
| L | `RFDETRLarge` | `rfdetr-large` | 75.1 | 56.5 | 6.8 | 33.9 | 704x704 | Apache 2.0 |
|
||||
| XL | `RFDETRXLarge` △ | `rfdetr-xlarge` | 77.4 | 58.6 | 11.5 | 126.4 | 700x700 | PML 1.0 |
|
||||
| 2XL | `RFDETR2XLarge` △ | `rfdetr-2xlarge` | 78.5 | 60.1 | 17.2 | 126.9 | 880x880 | PML 1.0 |
|
||||
|
||||
> △ Requires the `rfdetr_plus` extension: `pip install rfdetr[plus]`
|
||||
|
||||
## Run on an Image
|
||||
|
||||
Perform inference on an image using either the `rfdetr` package or the `inference` package. To use a different model size, select the corresponding class or alias from the table above.
|
||||
|
||||
=== "rfdetr"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
detections = model.predict("https://media.roboflow.com/dog.jpg", threshold=0.5)
|
||||
|
||||
labels = [f"{COCO_CLASSES[class_id]}" for class_id in detections.class_id]
|
||||
|
||||
annotated_image = sv.BoxAnnotator().annotate(detections.metadata["source_image"], detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
|
||||
```
|
||||
|
||||
=== "inference"
|
||||
|
||||
```python
|
||||
import requests
|
||||
import supervision as sv
|
||||
from PIL import Image
|
||||
from inference import get_model
|
||||
|
||||
model = get_model("rfdetr-medium")
|
||||
|
||||
image = Image.open(requests.get("https://media.roboflow.com/dog.jpg", stream=True).raw)
|
||||
predictions = model.infer(image, confidence=0.5)[0]
|
||||
detections = sv.Detections.from_inference(predictions)
|
||||
|
||||
annotated_image = sv.BoxAnnotator().annotate(image, detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections)
|
||||
```
|
||||
|
||||
!!! note "Using COCO classes vs. fine-tuned model classes"
|
||||
|
||||
`COCO_CLASSES` works for COCO-pretrained models (80 COCO classes, indexed 0-79).
|
||||
For fine-tuned models, use `detections.data["class_name"]` instead — it resolves
|
||||
class names from the checkpoint and works for both COCO and custom datasets.
|
||||
|
||||
For memory-constrained inference-only deployments with the `rfdetr` package, optimize the loaded model in place before
|
||||
calling `predict()`. Pass `dtype="float16"` to halve weight memory in addition to clearing the base model reference.
|
||||
This operation is irreversible — to restore the original model, create a new `RFDETR` instance:
|
||||
|
||||
```python
|
||||
model.optimize_for_inference(compile=False, inplace=True, dtype="float16")
|
||||
```
|
||||
|
||||
## Run on video, webcam, or RTSP stream
|
||||
|
||||
These examples use OpenCV for decoding and display. Replace `<SOURCE_VIDEO_PATH>`, `<WEBCAM_INDEX>`, and `<RTSP_STREAM_URL>` with your inputs. `<WEBCAM_INDEX>` is usually `0` for the default camera.
|
||||
|
||||
=== "video"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
video_capture = cv2.VideoCapture("<SOURCE_VIDEO_PATH>")
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError("Failed to open video source: <SOURCE_VIDEO_PATH>")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
detections = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]
|
||||
|
||||
annotated_frame = sv.BoxAnnotator().annotate(frame_bgr, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RF-DETR Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "webcam"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
WEBCAM_INDEX = 0
|
||||
video_capture = cv2.VideoCapture(WEBCAM_INDEX)
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError(f"Failed to open webcam: {WEBCAM_INDEX}")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
detections = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]
|
||||
|
||||
annotated_frame = sv.BoxAnnotator().annotate(frame_bgr, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RF-DETR Webcam", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "stream"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRMedium()
|
||||
|
||||
video_capture = cv2.VideoCapture("<RTSP_STREAM_URL>")
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError("Failed to open RTSP stream: <RTSP_STREAM_URL>")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
detections = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]
|
||||
|
||||
annotated_frame = sv.BoxAnnotator().annotate(frame_bgr, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RF-DETR RTSP", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
@@ -0,0 +1,203 @@
|
||||
---
|
||||
description: Run RF-DETR keypoint detection on images, video, and streams. COCO-pretrained preview model predicts 17 person keypoints with 71.8 AP at 9.7 ms on NVIDIA T4.
|
||||
---
|
||||
|
||||
# Run an RF-DETR Keypoint Model
|
||||
|
||||
RF-DETR Keypoint is a real-time transformer architecture for keypoint detection, built on a DINOv2 vision transformer backbone. The preview model is pretrained on the Microsoft COCO dataset and predicts 17 body keypoints per detected person.
|
||||
|
||||

|
||||
|
||||
!!! note "Preview model"
|
||||
|
||||
`RFDETRKeypointPreview` is an early-access release. Fine-tuning on custom keypoint datasets is the primary intended use case. See [Keypoint Preview Parameters](../train/training-parameters.md#keypoint-preview-parameters) for training configuration. API surface and checkpoint weights may change before the stable release.
|
||||
|
||||
## Pre-trained Checkpoints
|
||||
|
||||
RF-DETR Keypoint outperforms YOLO26-pose X and YOLO11-pose X at comparable latency on MS COCO. Latency measured on NVIDIA T4, TensorRT FP16, batch size 1.
|
||||
|
||||
{ width=560 }
|
||||
|
||||
| Model | RF-DETR package class | COCO AP<sub>50:95</sub> | Latency (ms) | Params (M) | Resolution | License |
|
||||
| :----------------: | :---------------------: | :---------------------: | :----------: | :--------: | :--------: | :--------: |
|
||||
| Keypoint (Preview) | `RFDETRKeypointPreview` | 71.8 | 9.7 | 126.4 | 576x576 | Apache 2.0 |
|
||||
|
||||
> The keypoint model is available in the `rfdetr` package only. It is not yet available via the `inference` package.
|
||||
|
||||
> Benchmark evaluated on COCO val2017 person keypoints (AP<sub>50:95</sub>) with the standard COCO 17-keypoint OKS sigmas; latency on NVIDIA T4, TensorRT FP16, batch size 1.
|
||||
|
||||
## Run on an Image
|
||||
|
||||
Perform inference on an image using the `rfdetr` package. `model.predict()` returns an [`sv.KeyPoints`](https://supervision.roboflow.com/latest/keypoint/core/) object containing skeleton coordinates and per-keypoint confidence scores for each detected person.
|
||||
|
||||
=== "rfdetr"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRKeypointPreview
|
||||
|
||||
model = RFDETRKeypointPreview()
|
||||
|
||||
image_bgr = cv2.imread("/path/to/image.jpg")
|
||||
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
|
||||
key_points = model.predict(image_rgb, threshold=0.5)
|
||||
|
||||
annotated_image = sv.VertexAnnotator().annotate(image_rgb, key_points)
|
||||
```
|
||||
|
||||

|
||||
|
||||
## Understanding the Output
|
||||
|
||||
`model.predict()` returns an `sv.KeyPoints` object. The fields most commonly used downstream:
|
||||
|
||||
| Field | Shape | Description |
|
||||
| --------------------------------- | -------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `key_points.xy` | `(N, K, 2)` | Pixel coordinates of each keypoint per detected instance |
|
||||
| `key_points.keypoint_confidence` | `(N, K)` | Per-keypoint findability score; use to filter low-confidence points |
|
||||
| `key_points.detection_confidence` | `(N,)` | Per-instance detection score; this is what `threshold` filters on. For keypoint models it includes the default uncertainty fusion term normalized to `[0, 1)`. |
|
||||
| `key_points.class_id` | `(N,)` | Model label ID for each detection. COCO-pretrained checkpoints use sparse COCO category IDs (1–90). Fine-tuned active-first keypoint checkpoints use normal 0-based class IDs; in the one-class preview setup, `class_id=0` is the foreground class and `class_id=1` is `"__background__"`. Legacy background-first keypoint checkpoints use slot 0 as `"__background__"` and start foreground classes at slot 1. Use `key_points.data["class_name"]` for name resolution rather than indexing your class list by `class_id`. |
|
||||
| `key_points.data["class_name"]` | `(N,)` | Class names resolved from `class_id`; prefer this over indexing a class-name list directly. |
|
||||
| `key_points.data["xyxy"]` | `(N, 4)` | Bounding box for each detected instance in `[x1, y1, x2, y2]` format |
|
||||
| `key_points.data["source_image"]` | list of arrays | Source frame stored once per detection; all N entries are the same array — use `[0]` to access it |
|
||||
|
||||
`K=17` for the pretrained COCO person-keypoint preview checkpoint. Fine-tuned checkpoints use the keypoint count from their dataset schema, so custom keypoint datasets can return any `K` supported by their COCO keypoint annotations.
|
||||
|
||||
Keypoints with `visible=False` are skipped by supervision annotators. To hide low-confidence joints manually, threshold `key_points.keypoint_confidence` and set matching entries to `False` in `key_points.visible`.
|
||||
|
||||
For fine-tuning on a custom keypoint dataset, see [Keypoint preview custom datasets](../train/index.md#keypoint-preview-custom-datasets).
|
||||
|
||||
## Run on video, webcam, or RTSP stream
|
||||
|
||||
These examples use OpenCV for decoding and display. Replace `<SOURCE_VIDEO_PATH>`, `<WEBCAM_INDEX>`, and `<RTSP_STREAM_URL>` with your inputs. `<WEBCAM_INDEX>` is usually `0` for the default camera.
|
||||
|
||||
=== "video"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRKeypointPreview
|
||||
|
||||
model = RFDETRKeypointPreview()
|
||||
|
||||
video_capture = cv2.VideoCapture("<SOURCE_VIDEO_PATH>")
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError("Failed to open video source: <SOURCE_VIDEO_PATH>")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
key_points = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
annotated_frame = sv.VertexAnnotator().annotate(frame_bgr, key_points)
|
||||
|
||||
cv2.imshow("RF-DETR Keypoint Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "webcam"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRKeypointPreview
|
||||
|
||||
model = RFDETRKeypointPreview()
|
||||
|
||||
WEBCAM_INDEX = 0 # Change this to the desired webcam index (e.g., 1, 2, ...)
|
||||
video_capture = cv2.VideoCapture(WEBCAM_INDEX)
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError(f"Failed to open webcam: {WEBCAM_INDEX}")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
key_points = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
annotated_frame = sv.VertexAnnotator().annotate(frame_bgr, key_points)
|
||||
|
||||
cv2.imshow("RF-DETR Keypoint Webcam", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "stream"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRKeypointPreview
|
||||
|
||||
model = RFDETRKeypointPreview()
|
||||
|
||||
video_capture = cv2.VideoCapture("<RTSP_STREAM_URL>")
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError("Failed to open RTSP stream: <RTSP_STREAM_URL>")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
key_points = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
annotated_frame = sv.VertexAnnotator().annotate(frame_bgr, key_points)
|
||||
|
||||
cv2.imshow("RF-DETR Keypoint RTSP", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
## Visualization
|
||||
|
||||
`supervision` provides several keypoint annotators. Choose based on what you want to draw.
|
||||
|
||||
=== "EdgeAnnotator"
|
||||
|
||||
Draws skeleton edges (lines between connected joints). Edges where either endpoint has `visible=False` are skipped automatically.
|
||||
|
||||
```python
|
||||
annotated = sv.EdgeAnnotator().annotate(image, key_points)
|
||||
```
|
||||
|
||||
=== "VertexAnnotator"
|
||||
|
||||
Draws a dot at each keypoint. Keypoints with `visible=False` are skipped automatically.
|
||||
|
||||
```python
|
||||
annotated = sv.VertexAnnotator().annotate(image, key_points)
|
||||
```
|
||||
|
||||
=== "VertexEllipseAnnotator"
|
||||
|
||||
Draws covariance ellipses from `key_points.data["covariance"]`, giving a visual footprint of per-keypoint uncertainty.
|
||||
|
||||
```python
|
||||
annotated = sv.VertexEllipseAnnotator().annotate(image, key_points)
|
||||
```
|
||||
|
||||
=== "VertexEllipseHaloAnnotator"
|
||||
|
||||
Draws the same covariance uncertainty with a soft halo for improved contrast on busy backgrounds.
|
||||
|
||||
```python
|
||||
annotated = sv.VertexEllipseHaloAnnotator().annotate(image, key_points)
|
||||
```
|
||||
@@ -0,0 +1,177 @@
|
||||
---
|
||||
description: Run RF-DETR instance segmentation on images, video, and streams. Mask predictions with 3.4-21.8 ms latency using DINOv2 backbone.
|
||||
---
|
||||
|
||||
# Run an RF-DETR Instance Segmentation Model
|
||||
|
||||
RF-DETR is a real-time transformer architecture for instance segmentation, built on a DINOv2 vision transformer backbone. The base models are trained on the Microsoft COCO dataset and achieve strong accuracy and latency trade-offs.
|
||||
|
||||
## Pre-trained Checkpoints
|
||||
|
||||
RF-DETR-Seg offers model sizes from Nano to 2XLarge, allowing trade-offs between accuracy, latency, and parameter count. All latency numbers were measured on an NVIDIA T4 using TensorRT, FP16, and batch size 1.
|
||||
|
||||
| Size | RF-DETR package class | Inference package alias | COCO AP<sub>50</sub> | COCO AP<sub>50:95</sub> | Latency (ms) | Params (M) | Resolution |
|
||||
| :--: | :-------------------: | :---------------------- | :------------------: | :---------------------: | :----------: | :--------: | :--------: |
|
||||
| N | `RFDETRSegNano` | `rfdetr-seg-nano` | 63.0 | 40.3 | 3.4 | 33.6 | 312x312 |
|
||||
| S | `RFDETRSegSmall` | `rfdetr-seg-small` | 66.2 | 43.1 | 4.4 | 33.7 | 384x384 |
|
||||
| M | `RFDETRSegMedium` | `rfdetr-seg-medium` | 68.4 | 45.3 | 5.9 | 35.7 | 432x432 |
|
||||
| L | `RFDETRSegLarge` | `rfdetr-seg-large` | 70.5 | 47.1 | 8.8 | 36.2 | 504x504 |
|
||||
| XL | `RFDETRSegXLarge` | `rfdetr-seg-xlarge` | 72.2 | 48.8 | 13.5 | 38.1 | 624x624 |
|
||||
| 2XL | `RFDETRSeg2XLarge` | `rfdetr-seg-2xlarge` | 73.1 | 49.9 | 21.8 | 38.6 | 768x768 |
|
||||
|
||||
## Run on an Image
|
||||
|
||||
Perform inference on an image using either the `rfdetr` package or the `inference` package. To use a different model size, select the corresponding class or alias from the table above.
|
||||
|
||||
=== "rfdetr"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRSegMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRSegMedium()
|
||||
|
||||
detections = model.predict("https://media.roboflow.com/dog.jpg", threshold=0.5)
|
||||
|
||||
labels = [f"{COCO_CLASSES[class_id]}" for class_id in detections.class_id]
|
||||
|
||||
annotated_image = sv.MaskAnnotator().annotate(detections.metadata["source_image"], detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
|
||||
```
|
||||
|
||||
=== "inference"
|
||||
|
||||
```python
|
||||
import requests
|
||||
import supervision as sv
|
||||
from PIL import Image
|
||||
from inference import get_model
|
||||
|
||||
model = get_model("rfdetr-seg-medium")
|
||||
|
||||
image = Image.open(requests.get("https://media.roboflow.com/dog.jpg", stream=True).raw)
|
||||
predictions = model.infer(image, confidence=0.5)[0]
|
||||
detections = sv.Detections.from_inference(predictions)
|
||||
|
||||
annotated_image = sv.MaskAnnotator().annotate(image, detections)
|
||||
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections)
|
||||
```
|
||||
|
||||
For memory-constrained inference-only deployments with the `rfdetr` package, optimize the loaded model in place before
|
||||
calling `predict()`. Pass `dtype="float16"` to halve weight memory in addition to clearing the base model reference.
|
||||
This operation is irreversible — to restore the original model, create a new `RFDETR` instance:
|
||||
|
||||
```python
|
||||
model.optimize_for_inference(compile=False, inplace=True, dtype="float16")
|
||||
```
|
||||
|
||||
## Run on video, webcam, or RTSP stream
|
||||
|
||||
These examples use OpenCV for decoding and display. Replace `<SOURCE_VIDEO_PATH>`, `<WEBCAM_INDEX>`, and `<RTSP_STREAM_URL>` with your inputs. `<WEBCAM_INDEX>` is usually `0` for the default camera.
|
||||
|
||||
=== "video"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRSegMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRSegMedium()
|
||||
|
||||
video_capture = cv2.VideoCapture("<SOURCE_VIDEO_PATH>")
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError("Failed to open video source: <SOURCE_VIDEO_PATH>")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
detections = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]
|
||||
|
||||
annotated_frame = sv.MaskAnnotator().annotate(frame_bgr, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RF-DETR-Seg Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "webcam"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRSegMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRSegMedium()
|
||||
|
||||
WEBCAM_INDEX = 0 # Change this to the desired webcam index (e.g., 1, 2, ...)
|
||||
video_capture = cv2.VideoCapture(WEBCAM_INDEX)
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError(f"Failed to open webcam: {WEBCAM_INDEX}")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
detections = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]
|
||||
|
||||
annotated_frame = sv.MaskAnnotator().annotate(frame_bgr, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RF-DETR-Seg Webcam", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "stream"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from rfdetr import RFDETRSegMedium
|
||||
from rfdetr.assets.coco_classes import COCO_CLASSES
|
||||
|
||||
model = RFDETRSegMedium()
|
||||
|
||||
video_capture = cv2.VideoCapture("<RTSP_STREAM_URL>")
|
||||
if not video_capture.isOpened():
|
||||
raise RuntimeError("Failed to open RTSP stream: <RTSP_STREAM_URL>")
|
||||
|
||||
while True:
|
||||
success, frame_bgr = video_capture.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
detections = model.predict(frame_rgb, threshold=0.5)
|
||||
|
||||
labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]
|
||||
|
||||
annotated_frame = sv.MaskAnnotator().annotate(frame_bgr, detections)
|
||||
annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)
|
||||
|
||||
cv2.imshow("RF-DETR-Seg RTSP", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
video_capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
@@ -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. |
|
||||
Reference in New Issue
Block a user