Custom optimizer and LR scheduler in RF-DETR¶

RF-DETR fine-tunes with a managed AdamW optimizer and a step LR schedule by default. Both are swappable through TrainConfig without touching the training loop — you can select any torch.optim optimizer, any torch.optim.lr_scheduler schedule, a third-party optimizer (e.g. pytorch-optimizer) by import path, or a functools.partial with your own hyperparameters baked in.

This cookbook fine-tunes RFDETRSmall on a small Roboflow Universe dataset while replacing both the optimizer and the scheduler, then verifies the swap actually took effect on the built LightningModule.

The three selection modes¶

TrainConfig.optimizer and TrainConfig.lr_scheduler each accept one of:

Mode	`optimizer` example	`lr_scheduler` example	Who supplies hyperparameters
Managed short name	`"adamw"`, `"sgd"`	`"step"`, `"cosine"`	RF-DETR injects `lr` / `weight_decay`; presets own warmup + full-run sizing
Explicit import path	`"torch.optim.AdamW"`, `"pytorch_optimizer.Lion"`	`"torch.optim.lr_scheduler.CosineAnnealingLR"`	You, via `optimizer_kwargs` / `lr_scheduler_kwargs` (RF-DETR injects nothing)
Callable / `functools.partial`	`functools.partial(torch.optim.AdamW, weight_decay=1e-4)`	`functools.partial(StepLR, step_size=30)`	Baked into the callable; the `*_kwargs` field is ignored (with a warning)

See Training customization → Custom optimizer / LR scheduler and Training parameters for the full reference.

Setup¶

Install rfdetr with the train extra (PyTorch Lightning, COCO eval) plus roboflow for the dataset download. A GPU is recommended; this small dataset also runs on CPU for a quick smoke test.

In [ ]:

Copied!

!pip install -q "rfdetr[train]>=1.9.0" roboflow
!pip install -q "rfdetr[train]>=1.9.0" roboflow

In [ ]:

Copied!





"""Fine-tune RF-DETR with a custom optimizer and LR scheduler on a Roboflow Universe dataset."""

import json
import os
from pathlib import Path
from typing import Any

from rfdetr import RFDETRSmall
from rfdetr.config import TrainConfig
from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
from rfdetr.utilities.reproducibility import seed_all
"""Fine-tune RF-DETR with a custom optimizer and LR scheduler on a Roboflow Universe dataset."""

import json
import os
from pathlib import Path
from typing import Any

from rfdetr import RFDETRSmall
from rfdetr.config import TrainConfig
from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
from rfdetr.utilities.reproducibility import seed_all

In [ ]:

Copied!





PROJECT_ROOT = Path(__file__).resolve().parent if "__file__" in globals() else Path.cwd()
DATASETS_DIR = PROJECT_ROOT / "datasets"

# A small, single-domain Roboflow Universe COCO dataset keeps this demo fast — the point here is the
# optimizer / scheduler wiring, not the dataset. Swap in any COCO-format Universe export you like.
DATASET_INFO: dict[str, Any] = {
    "name": "Football Player Detection",
    "workspace": "football-gozni",
    "project": "football-player-detection-bfswn",
    "version": 1,
    "source_url": "https://universe.roboflow.com/football-gozni/football-player-detection-bfswn",
}
DATASET_KEY = "football"

OUTPUT_DIR = PROJECT_ROOT / "output" / "custom_optimizer_scheduler"

SEED = 7
EPOCHS = 20
BATCH_SIZE = 8
GRAD_ACCUM_STEPS = 2
NUM_WORKERS = 8
RESOLUTION = 512
LR = 1e-4
LR_ENCODER = 1e-4
PROJECT_ROOT = Path(__file__).resolve().parent if "__file__" in globals() else Path.cwd()
DATASETS_DIR = PROJECT_ROOT / "datasets"

# A small, single-domain Roboflow Universe COCO dataset keeps this demo fast — the point here is the
# optimizer / scheduler wiring, not the dataset. Swap in any COCO-format Universe export you like.
DATASET_INFO: dict[str, Any] = {
    "name": "Football Player Detection",
    "workspace": "football-gozni",
    "project": "football-player-detection-bfswn",
    "version": 1,
    "source_url": "https://universe.roboflow.com/football-gozni/football-player-detection-bfswn",
}
DATASET_KEY = "football"

OUTPUT_DIR = PROJECT_ROOT / "output" / "custom_optimizer_scheduler"

SEED = 7
EPOCHS = 20
BATCH_SIZE = 8
GRAD_ACCUM_STEPS = 2
NUM_WORKERS = 8
RESOLUTION = 512
LR = 1e-4
LR_ENCODER = 1e-4

1 - Download the dataset¶

Get a free Roboflow API key at app.roboflow.com/settings/api and expose it as ROBOFLOW_API_KEY (an environment variable locally, or a Colab secret with the same name). The download is idempotent — re-running skips the transfer if the target directory already exists.

In [ ]:

Copied!





try:
    from google.colab import userdata

    try:
        ROBOFLOW_API_KEY = userdata.get("ROBOFLOW_API_KEY") or ""
    except Exception:
        ROBOFLOW_API_KEY = ""
except ImportError:
    ROBOFLOW_API_KEY = ""

if not ROBOFLOW_API_KEY:
    ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY", "")
if not ROBOFLOW_API_KEY:
    raise RuntimeError(
        "ROBOFLOW_API_KEY not found. "
        "In Colab: add it via Secrets (key icon). "
        "Locally: set the environment variable before running."
    )

from roboflow import Roboflow

rf = Roboflow(api_key=ROBOFLOW_API_KEY)
dataset = (
    rf.workspace(str(DATASET_INFO["workspace"]))
    .project(str(DATASET_INFO["project"]))
    .version(int(DATASET_INFO["version"]))
    .download("coco", location=str(DATASETS_DIR / DATASET_KEY))
)
DATASET_DIR = Path(dataset.location)
TRAIN_ANNOTATIONS = DATASET_DIR / "train" / "_annotations.coco.json"
print(f"dataset_dir={DATASET_DIR}")
try:
    from google.colab import userdata

    try:
        ROBOFLOW_API_KEY = userdata.get("ROBOFLOW_API_KEY") or ""
    except Exception:
        ROBOFLOW_API_KEY = ""
except ImportError:
    ROBOFLOW_API_KEY = ""

if not ROBOFLOW_API_KEY:
    ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY", "")
if not ROBOFLOW_API_KEY:
    raise RuntimeError(
        "ROBOFLOW_API_KEY not found. "
        "In Colab: add it via Secrets (key icon). "
        "Locally: set the environment variable before running."
    )

from roboflow import Roboflow

rf = Roboflow(api_key=ROBOFLOW_API_KEY)
dataset = (
    rf.workspace(str(DATASET_INFO["workspace"]))
    .project(str(DATASET_INFO["project"]))
    .version(int(DATASET_INFO["version"]))
    .download("coco", location=str(DATASETS_DIR / DATASET_KEY))
)
DATASET_DIR = Path(dataset.location)
TRAIN_ANNOTATIONS = DATASET_DIR / "train" / "_annotations.coco.json"
print(f"dataset_dir={DATASET_DIR}")

2 - Infer class names¶

RF-DETR uses 0-based class indices; sort the COCO categories by their original id and read names in order.

In [ ]:

Copied!





with TRAIN_ANNOTATIONS.open() as _f:
    _coco = json.load(_f)

CLASS_NAMES: list[str] = [cat["name"] for cat in sorted(_coco["categories"], key=lambda c: c["id"])]
NUM_CLASSES = len(CLASS_NAMES)
print(f"class_names={CLASS_NAMES}")
print(f"num_classes={NUM_CLASSES}")
with TRAIN_ANNOTATIONS.open() as _f:
    _coco = json.load(_f)

CLASS_NAMES: list[str] = [cat["name"] for cat in sorted(_coco["categories"], key=lambda c: c["id"])]
NUM_CLASSES = len(CLASS_NAMES)
print(f"class_names={CLASS_NAMES}")
print(f"num_classes={NUM_CLASSES}")

3 - Configure a custom optimizer and LR scheduler¶

This run replaces both defaults using explicit import paths:

Optimizer — "torch.optim.AdamW". In explicit-path mode RF-DETR builds the class from its own parameter groups (which already carry the per-group lr / lr_encoder) plus optimizer_kwargs only; it injects no lr or weight_decay, so anything the optimizer needs goes in optimizer_kwargs.
Scheduler — "torch.optim.lr_scheduler.CosineAnnealingLR" with T_max=EPOCHS, stepped once per epoch (lr_scheduler_interval="epoch"). RF-DETR injects no T_max, so it is passed explicitly.
Warmup — with warmup_epochs > 0, an explicit scheduler is automatically prepended with a linear warmup ramp via SequentialLR (managed presets bake warmup into their own schedule instead).

Two alternative ways to express the same intent (uncomment to try):

# Managed short name — RF-DETR injects lr / weight_decay for you:
optimizer="sgd", optimizer_kwargs={"momentum": 0.9, "nesterov": True}

# Callable / functools.partial — bake args in; optimizer_kwargs is ignored:
import functools
optimizer=functools.partial(torch.optim.AdamW, weight_decay=1e-4)
lr_scheduler=functools.partial(torch.optim.lr_scheduler.StepLR, step_size=10, gamma=0.1)

In [ ]:

Copied!





seed_all(SEED)
variant = RFDETRSmall(  # type: ignore[no-untyped-call]
    num_classes=NUM_CLASSES,
    resolution=RESOLUTION,
)
variant.model_config.model_name = type(variant).__name__

train_config = TrainConfig(
    dataset_file="roboflow",
    dataset_dir=str(DATASET_DIR),
    output_dir=str(OUTPUT_DIR),
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    grad_accum_steps=GRAD_ACCUM_STEPS,
    num_workers=NUM_WORKERS,
    lr=LR,
    lr_encoder=LR_ENCODER,
    warmup_epochs=1,
    # --- custom optimizer (explicit import path) ---
    optimizer="torch.optim.AdamW",
    optimizer_kwargs={"betas": (0.9, 0.999), "weight_decay": 1e-4},
    # --- custom LR scheduler (explicit import path) ---
    lr_scheduler="torch.optim.lr_scheduler.CosineAnnealingLR",
    lr_scheduler_kwargs={"T_max": EPOCHS},
    lr_scheduler_interval="epoch",
    use_ema=False,
    run_test=False,
    multi_scale=False,
    expanded_scales=False,
    tensorboard=False,
    wandb=False,
    mlflow=False,
    clearml=False,
    class_names=CLASS_NAMES,
    progress_bar="tqdm",
)

datamodule = RFDETRDataModule(variant.model_config, train_config)
model = RFDETRModelModule(variant.model_config, train_config)
trainer = build_trainer(train_config, variant.model_config)
seed_all(SEED)
variant = RFDETRSmall(  # type: ignore[no-untyped-call]
    num_classes=NUM_CLASSES,
    resolution=RESOLUTION,
)
variant.model_config.model_name = type(variant).__name__

train_config = TrainConfig(
    dataset_file="roboflow",
    dataset_dir=str(DATASET_DIR),
    output_dir=str(OUTPUT_DIR),
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    grad_accum_steps=GRAD_ACCUM_STEPS,
    num_workers=NUM_WORKERS,
    lr=LR,
    lr_encoder=LR_ENCODER,
    warmup_epochs=1,
    # --- custom optimizer (explicit import path) ---
    optimizer="torch.optim.AdamW",
    optimizer_kwargs={"betas": (0.9, 0.999), "weight_decay": 1e-4},
    # --- custom LR scheduler (explicit import path) ---
    lr_scheduler="torch.optim.lr_scheduler.CosineAnnealingLR",
    lr_scheduler_kwargs={"T_max": EPOCHS},
    lr_scheduler_interval="epoch",
    use_ema=False,
    run_test=False,
    multi_scale=False,
    expanded_scales=False,
    tensorboard=False,
    wandb=False,
    mlflow=False,
    clearml=False,
    class_names=CLASS_NAMES,
    progress_bar="tqdm",
)

datamodule = RFDETRDataModule(variant.model_config, train_config)
model = RFDETRModelModule(variant.model_config, train_config)
trainer = build_trainer(train_config, variant.model_config)

4 - Fine-tune¶

trainer.fit runs the loop with the custom optimizer + scheduler. RFDETRModelModule.configure_optimizers() builds them from the config during setup — it needs the connected trainer to size the schedule, so it runs inside fit, not standalone. On this small dataset 20 epochs is enough to see the cosine LR decay take effect; scale EPOCHS up for a real run.

In [ ]:

Copied!

trainer.fit(model, datamodule=datamodule)
trainer.fit(model, datamodule=datamodule)

5 - Verify the swap took effect¶

After fit, PyTorch Lightning exposes the objects it actually trained with. Confirm the custom optimizer and scheduler were selected: an AdamW optimizer and a SequentialLR (the outer wrapper is the automatic linear-warmup prepended because warmup_epochs > 0; its second phase is the CosineAnnealingLR).

In [ ]:

Copied!





_optimizer = trainer.optimizers[0]
_scheduler = trainer.lr_scheduler_configs[0].scheduler
print(f"optimizer = {type(_optimizer).__name__}")
print(f"scheduler = {type(_scheduler).__name__}")
print(f"param_group_lrs = {[round(g['lr'], 6) for g in _optimizer.param_groups]}")
_optimizer = trainer.optimizers[0]
_scheduler = trainer.lr_scheduler_configs[0].scheduler
print(f"optimizer = {type(_optimizer).__name__}")
print(f"scheduler = {type(_scheduler).__name__}")
print(f"param_group_lrs = {[round(g['lr'], 6) for g in _optimizer.param_groups]}")

Where to go next¶

Swap optimizer="torch.optim.AdamW" for a third-party optimizer by import path, e.g. optimizer="pytorch_optimizer.Lion" (install pytorch-optimizer first).
Try a different schedule — "torch.optim.lr_scheduler.OneCycleLR" (needs total_steps in lr_scheduler_kwargs) or the managed "cosine" preset (lr_scheduler_kwargs={"min_factor": 0.1}).
Wrapper optimizers (SAM, Lookahead) need a custom configure_optimizers — see the customization guide for the caveats.