# How to Specify Augmentation Backends for RF-DETR Training

> Learn how to specify augmentation backends for RF-DETR training. Choose from Albumentations, Torchvision, or Kornia via the augmentation_backend parameter to boost preprocessing performance.

- Repository: [Roboflow/rf-detr](https://github.com/roboflow/rf-detr)
- Tags: how-to-guide
- Published: 2026-09-08

---

**RF-DETR supports three distinct augmentation pipelines—CPU-based Albumentations, CPU-based Torchvision, and GPU-based Kornia—that you select via the `augmentation_backend` parameter in `TrainConfig` to optimize preprocessing performance for your hardware.**

RF-DETR provides flexible data augmentation through pluggable backends that minimize data transfer overhead between CPU and GPU. By configuring the `augmentation_backend` field in `TrainConfig`, you control whether image transformations run on the host CPU before batching or directly on GPU tensors during training.

## Available Augmentation Backends

RF-DETR implements three augmentation strategies, each optimized for different environment constraints and performance requirements.

### CPU Albumentations Backend

The **Albumentations** backend runs on the host CPU using the `albumentations` library. This is the default pipeline when you specify `"cpu"` or `"auto"` and the library is installed. It applies geometric and photometric transforms to numpy arrays before tensors are moved to the GPU. If Albumentations is not installed, the system automatically falls back to the Torchvision CPU pipeline.

### Torchvision CPU Backend

The **Torchvision** backend provides a pure-CPU pipeline using only `torchvision` transforms. Select this by setting `augmentation_backend="torchvision"` when you require deterministic behavior independent of optional packages, or when Albumentations is unavailable in your environment.

### GPU Kornia Backend

The **Kornia** backend executes augmentations directly on GPU tensors, eliminating the CPU-GPU round-trip for random transforms. Specify `"kornia"` or `"gpu"` to enable this pipeline. This backend requires a CUDA-capable device and the optional `kfdetr[augment]` dependency package. If CUDA is unavailable or Kornia is not installed, the training initialization raises a `RuntimeError` or `ImportError`.

## Configuring the Backend in TrainConfig

The augmentation backend is defined in the `TrainConfig` class located in [`src/rfdetr/config.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/config.py). The field accepts both enum values and string literals:

```python

# src/rfdetr/config.py

class TrainConfig(BaseConfig):
    ...
    augmentation_backend: AugmentationBackend | Literal["cpu", "auto"] = "cpu"

```

Valid string values include:
- `"cpu"` or `"auto"` – Defaults to Albumentations if available, otherwise Torchvision
- `"albumentations"` or `"albu"` – Forces Albumentations (raises error if missing)
- `"torchvision"` or `"tv"` – Forces Torchvision transforms
- `"kornia"` or `"gpu"` – Forces GPU-side Kornia transforms

## Backend Resolution and Validation

When training starts, `RFDETRTrainingModule.setup()` in [`src/rfdetr/training/module_data.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/training/module_data.py) reads your configuration and resolves the string to a concrete `AugmentationBackend` enum:

```python

# src/rfdetr/training/module_data.py

requested_backend = self.train_config.augmentation_backend
resolved = resolve_augmentation_backend(requested_backend,
                                        has_cuda=_has_cuda_device())
self._resolved_augmentation_backend = resolved

```

The resolution logic resides in [`src/rfdetr/datasets/kornia_transforms.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/datasets/kornia_transforms.py). The `resolve_augmentation_backend` function maps inputs to enum values while performing safety checks:

```python

# src/rfdetr/datasets/kornia_transforms.py

def resolve_augmentation_backend(backend: str, *, has_cuda: bool | None = None) -> AugmentationBackend:
    if backend in (AugmentationBackend.ALBU, "albumentations", "albu"):
        _require_albu()
    if has_cuda is None:
        has_cuda = _has_cuda_device()
    return AugmentationBackend.from_str(backend, has_cuda=has_cuda)

```

This function validates that Albumentations is installed for ALBU requests and confirms CUDA availability for Kornia requests before returning the resolved enum.

## Dataset Pipeline Integration

Dataset builders utilize the resolved backend to determine whether to apply CPU-side augmentation or defer transforms to the GPU stage. In [`src/rfdetr/datasets/coco.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/datasets/coco.py), the builder checks the resolved backend to configure post-processing:

```python

# src/rfdetr/datasets/coco.py

augmentation_backend = getattr(args, "augmentation_backend", "cpu")
resolved_augmentation_backend = resolve_backend_for_build(augmentation_backend)
gpu_postprocess = is_gpu_postprocess(resolved_augmentation_backend)

```

When `gpu_postprocess` evaluates to `True` (indicating the Kornia backend), the CPU augmentation step is omitted and normalization operations occur on-device, reducing host-device synchronization overhead.

## Practical Configuration Examples

### Default Albumentations Pipeline

Use the default CPU backend with automatic package detection:

```python
from rfdetr import RFDETR
from rfdetr.config import TrainConfig

rf = RFDETR.train(
    TrainConfig(
        dataset_dir="my_dataset",
        augmentation_backend="cpu",  # Optional; this is the default

    )
)

```

### Force Torchvision Backend

Ensure deterministic CPU behavior without optional dependencies:

```python
rf = RFDETR.train(
    TrainConfig(
        dataset_dir="my_dataset",
        augmentation_backend="torchvision",
    )
)

```

### Enable GPU Kornia Acceleration

Require CUDA and the Kornia package for on-device augmentation:

```python
rf = RFDETR.train(
    TrainConfig(
        dataset_dir="my_dataset",
        augmentation_backend="kornia",  # or "gpu"

    )
)

```

### Automatic Backend Selection

Let the library select the best available backend based on installed packages and hardware:

```python
rf = RFDETR.train(
    TrainConfig(
        dataset_dir="my_dataset",
        augmentation_backend="auto",
    )
)

```

## Summary

- RF-DETR offers three augmentation pipelines: **Albumentations (CPU)**, **Torchvision (CPU)**, and **Kornia (GPU)** configurable via `TrainConfig`.
- Set the `augmentation_backend` parameter to `"cpu"`, `"torchvision"`, `"kornia"`, or `"auto"` to control preprocessing behavior.
- The resolution logic in [`src/rfdetr/datasets/kornia_transforms.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/datasets/kornia_transforms.py) validates environment compatibility and CUDA availability during training initialization.
- Dataset builders in [`src/rfdetr/datasets/coco.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/datasets/coco.py) use the resolved backend to skip CPU augmentation when GPU post-processing is active.

## Frequently Asked Questions

### What is the default augmentation backend in RF-DETR?

The default value is `"cpu"`, which resolves to Albumentations if the package is installed, otherwise falling back to Torchvision transforms. You can verify the resolved backend by inspecting `self._resolved_augmentation_backend` in the training module after initialization.

### What happens if I specify "kornia" but have no GPU?

The `resolve_augmentation_backend` function in [`src/rfdetr/datasets/kornia_transforms.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/datasets/kornia_transforms.py) detects CUDA availability during resolution. If `has_cuda` is `False` and you request `"kornia"` or `"gpu"`, the function raises a `RuntimeError` indicating that GPU augmentation requires a CUDA device.

### Can I use RF-DETR augmentation without installing Albumentations?

Yes. If you set `augmentation_backend="torchvision"` or rely on the default `"cpu"`/`"auto"` behavior when Albumentations is absent, the system automatically falls back to the Torchvision CPU pipeline. No error occurs unless you explicitly force `"albumentations"`.

### Where is the augmentation backend resolved during training initialization?

The resolution occurs in `RFDETRTrainingModule.setup()` within [`src/rfdetr/training/module_data.py`](https://github.com/roboflow/rf-detr/blob/main/src/rfdetr/training/module_data.py). This method calls `resolve_augmentation_backend` to convert your string configuration into a concrete enum value, storing the result in `self._resolved_augmentation_backend` for use by dataset loaders.