Train RF-DETR Nano on COCO2017 (high-end GPU)¶

Downloads the full public COCO2017 dataset and trains RF-DETR Nano for 40 epochs, end to end. This notebook is sized for a high-end GPU — an A100, an H100, or an RTX PRO 6000 (e.g. Colab Pro+ with an A100 runtime, or a GCP G4 VM) — and its batch size, RAM-disk step, and worker count all assume that class of machine. It still runs on a smaller GPU (T4/L4): switch batch_size back to "auto" in the training cell and the rest adapts, but expect the run to take far longer, and skip the RAM-disk move on a low-RAM host — it needs ~19 GB free in /dev/shm.

The throughput work is all in the shipped defaults — the notebook adds host-adaptive knobs (num_workers, seed), a batch size measured on high-end hardware, and a learning-rate schedule sized for a 40-epoch run, each explained where it is set.

Kept deliberately minimal: download, train, plot the metrics CSVLogger wrote during training. For dataset preview, checkpoint saving, and inference visualization, see fine-tune_detection.ipynb.

Why this notebook is simple on purpose¶

Recent releases moved several training-throughput fixes directly into the default configuration, so a stock run picks them up automatically:

• Validation forwards one model per epoch, not two. The base-model forward pass used to run alongside the EMA forward every validation batch; it is now skipped when use_ema=True (the default).
• grad_accum_steps now defaults to 1 instead of 4. Gradient accumulation is an explicit opt-in — raise batch_size for your GPU first, and reach for accumulation only when memory forces a smaller physical batch. Measured on one L4: batch_size=16, grad_accum_steps=1 ran 27% faster per epoch than batch_size=4, grad_accum_steps=4 at the same nominal effective batch, with equal mAP.
• eval_batch_size decouples the validation/test dataloaders from the training micro-batch size, so a small training batch no longer forces small (slower) evaluation batches.
• COCO mAP computation reads each image's detection scores once instead of once per detection.
• The transformer skips materializing tensors it would otherwise reuse unchanged on the single-feature-level path that Nano (and every current detection size) uses by default.
• Pre-training sanity-check validation is skipped by default.

None of this needs a flag — it is what RFDETRNano().train() already does. TensorBoard logging is turned off so the notebook does not require the loggers extra; CSVLogger is always on and is what section 5 plots.

Scope and honest expectations¶

COCO2017 has 118,287 training images and 5,000 validation images — two to three orders of magnitude larger than the small Roboflow Universe datasets in the other fine-tuning cookbooks. There is no COCO2017-scale timing measurement in this repository yet, so no fixed "N minutes per epoch" number is given here. Time your own first epoch (the progress bar prints per-epoch elapsed time) before committing GPU time to all 40. As a rough anchor: an internal L4 measurement on a ~6,900-image dataset with rfdetr-small at batch 16 took ~324 s/epoch (train + validation); COCO2017 has about 17x more training images and 2.5x more validation images, and rfdetr-nano is smaller and runs at a lower resolution than rfdetr-small, which pulls the other way. Expect the full 40-epoch run to take multiple hours even on a fast GPU.

On the target hardware (A100/H100/RTX PRO 6000) the full 40-epoch run fits comfortably in one session, and Colab Pro+ background execution keeps it alive with the browser closed. Training still checkpoints every checkpoint_interval epochs (10 by default) to last.ckpt as a safety net: if a session does disconnect, reopen the notebook, set train_config.resume to that file's path, and re-run the training cell to continue. On a free-tier T4/L4 the session is likely to end before 40 epochs do — plan on resuming.

COCO2017 also needs about 19 GB of disk once downloaded (the archives are deleted right after extraction to avoid a ~38 GB peak). Confirm your Colab runtime has that much free space before starting the download.

1 - Download COCO2017¶

The download comes first because it is by far the longest step and needs nothing installed — only wget and unzip, both already on a Colab runtime. Anything the environment setup below does to the session (including a runtime restart, if pip asks for one) leaves the extracted dataset on disk untouched.

Plain wget/unzip into the standard COCO layout RF-DETR's dataset_file="coco" loader expects: coco2017/train2017/, coco2017/val2017/, coco2017/annotations/instances_{train,val}2017.json.

Each archive is downloaded, extracted, and deleted before the next one starts, so the disk never holds more than one archive alongside the extracted data. The && chaining is deliberate: a notebook shell cell does not stop on a failed command, so a bare sequence would run unzip on a truncated download and rm on a failed extraction, then scroll past — the failure would only surface much later, as a missing-image error part-way into training. wget -c resumes an interrupted download instead of restarting it, and unzip -o overwrites files already extracted, so re-running this cell after a disconnected Colab session replaces any partially extracted file. The closing df -h shows the disk headroom left; the extracted dataset needs about 19 GB.

In [ ]:

!mkdir -p datasets/coco2017
!cd datasets/coco2017 && wget -c -q --show-progress http://images.cocodataset.org/zips/train2017.zip && unzip -o -q train2017.zip && rm -f train2017.zip
!cd datasets/coco2017 && wget -c -q --show-progress http://images.cocodataset.org/zips/val2017.zip && unzip -o -q val2017.zip && rm -f val2017.zip
!cd datasets/coco2017 && wget -c -q --show-progress http://images.cocodataset.org/annotations/annotations_trainval2017.zip && unzip -o -q annotations_trainval2017.zip && rm -f annotations_trainval2017.zip
!df -h datasets/coco2017

!mkdir -p datasets/coco2017 !cd datasets/coco2017 && wget -c -q --show-progress http://images.cocodataset.org/zips/train2017.zip && unzip -o -q train2017.zip && rm -f train2017.zip !cd datasets/coco2017 && wget -c -q --show-progress http://images.cocodataset.org/zips/val2017.zip && unzip -o -q val2017.zip && rm -f val2017.zip !cd datasets/coco2017 && wget -c -q --show-progress http://images.cocodataset.org/annotations/annotations_trainval2017.zip && unzip -o -q annotations_trainval2017.zip && rm -f annotations_trainval2017.zip !df -h datasets/coco2017

Verify the extraction before spending GPU time on it. COCO2017 ships 118,287 training and 5,000 validation images; a short count catches a truncated download or an out-of-disk extraction here, at the cost of a few seconds, instead of several epochs later when the training loop first reaches a missing file.

In [ ]:

from pathlib import Path

COCO_ROOT = Path("datasets/coco2017")
EXPECTED_IMAGES = {"train2017": 118_287, "val2017": 5_000}

for split, expected in EXPECTED_IMAGES.items():
    found = sum(1 for _ in (COCO_ROOT / split).glob("*.jpg"))
    assert found == expected, (
        f"{COCO_ROOT / split} holds {found} images, expected {expected}. "
        "Re-run the download cell above; check the output of `df -h` for a full disk."
    )
    annotations = COCO_ROOT / "annotations" / f"instances_{split}.json"
    assert annotations.is_file(), f"Missing {annotations}. Re-run the download cell above."

print("COCO2017 is complete.")

from pathlib import Path COCO_ROOT = Path("datasets/coco2017") EXPECTED_IMAGES = {"train2017": 118_287, "val2017": 5_000} for split, expected in EXPECTED_IMAGES.items(): found = sum(1 for _ in (COCO_ROOT / split).glob("*.jpg")) assert found == expected, ( f"{COCO_ROOT / split} holds {found} images, expected {expected}. " "Re-run the download cell above; check the output of `df -h` for a full disk." ) annotations = COCO_ROOT / "annotations" / f"instances_{split}.json" assert annotations.is_file(), f"Missing {annotations}. Re-run the download cell above." print("COCO2017 is complete.")

2 - Set up the environment¶

The training-throughput work described above has not shipped in a tagged release yet, so this cell installs from the develop branch. Once a release containing it is out, replace this with pip install -q "rfdetr[train,augment,visual]>=1.10.0".

GPU required. Training needs a CUDA GPU, and this notebook's batch size assumes a high-end one — in Colab: Runtime → Change runtime type → A100 GPU (Pro/Pro+ plans; also enable High-RAM so the /dev/shm move below fits).

In [ ]:

!pip install -q "rfdetr[train,augment,visual] @ git+https://github.com/roboflow/rf-detr.git@develop"

!pip install -q "rfdetr[train,augment,visual] @ git+https://github.com/roboflow/rf-detr.git@develop"

3 - Move the dataset to /dev/shm¶

/dev/shm is Linux's tmpfs RAM disk — after the move, every DataLoader read is a memory access with no storage I/O at all. On the high-end hosts this notebook targets, RAM is plentiful (a GCP G4 shape carries 180+ GB, an A100 Colab runtime 80+ GB; /dev/shm is sized to half of RAM by default), so the ~19 GB dataset fits. It is a real win there because cloud persistent disks throttle throughput in proportion to provisioned size — a typical 100-200 GB boot disk serves 118k random-access JPEG reads per epoch slowly. JPEG decode cost is unaffected by where the bytes come from; that CPU cost is what the high num_workers below parallelizes away.

A move (not a copy) keeps a single instance of the dataset on the machine. Two consequences to know:

• tmpfs does not survive a runtime restart — which is why this cell sits after the pip install (pip can restart the runtime). If the session restarts later anyway, the dataset is gone: re-run the download cell.
• If /dev/shm is too small, mv fails part-way and leaves files split across both locations. Recover with mv /dev/shm/coco2017/* datasets/coco2017/ and train from disk — or better, avoid it: on Colab pick a High-RAM runtime; df -h /dev/shm shows capacity before you commit.

The test -d guard makes the cell safe to re-run: once the dataset already sits in /dev/shm, the mv is skipped instead of nesting a second copy inside it.

In [ ]:

!df -h /dev/shm
!test -d /dev/shm/coco2017 || mv datasets/coco2017 /dev/shm/coco2017

!df -h /dev/shm !test -d /dev/shm/coco2017 || mv datasets/coco2017 /dev/shm/coco2017

4 - Load the model¶

RFDETRNano() loads the released COCO-pretrained Nano checkpoint by default — this notebook continues training from it rather than from random initialization. Training from scratch is not the recommended path for RF-DETR; continuing from the pretrained checkpoint on the same dataset it was pretrained on still exercises the full 40-epoch training loop and every optimization listed above, which is what this notebook demonstrates.

num_classes is left at its default of 90 — the standard COCO category-id space RF-DETR's dataset_file="coco" loader and the pretrained checkpoint both already use, so it does not need to be passed explicitly.

In [ ]:

import os
from pathlib import Path

import torch

from rfdetr import RFDETRNano
from rfdetr.visualize.training import plot_loss_metrics, plot_map_metrics

if not torch.cuda.is_available():
    raise RuntimeError("This notebook requires a CUDA GPU. In Colab: Runtime -> Change runtime type -> A100 GPU.")

# Repeated from the verification cell so this section still runs on its own after a runtime restart.
COCO_ROOT = Path("datasets/coco2017")
RAMDISK_ROOT = Path("/dev/shm/coco2017")
if RAMDISK_ROOT.is_dir():
    COCO_ROOT = RAMDISK_ROOT
OUTPUT_DIR = "output/det_coco2017_nano"
EPOCHS = 40
# Capped: measured on a 48-core host, uncapped workers sat ~15% CPU-utilized while holding ~70% of RAM.
NUM_WORKERS = min(os.cpu_count() or 2, 16)
# Batch 128 was measured at ~80% GPU memory on a high-end (80+ GB) GPU — sized for an A100, H100, or
# RTX PRO 6000. On a smaller GPU replace this with "auto" and the probe finds the largest batch that fits.
BATCH_SIZE = 128

model = RFDETRNano()  # type: ignore[no-untyped-call]

import os from pathlib import Path import torch from rfdetr import RFDETRNano from rfdetr.visualize.training import plot_loss_metrics, plot_map_metrics if not torch.cuda.is_available(): raise RuntimeError("This notebook requires a CUDA GPU. In Colab: Runtime -> Change runtime type -> A100 GPU.") # Repeated from the verification cell so this section still runs on its own after a runtime restart. COCO_ROOT = Path("datasets/coco2017") RAMDISK_ROOT = Path("/dev/shm/coco2017") if RAMDISK_ROOT.is_dir(): COCO_ROOT = RAMDISK_ROOT OUTPUT_DIR = "output/det_coco2017_nano" EPOCHS = 40 # Capped: measured on a 48-core host, uncapped workers sat ~15% CPU-utilized while holding ~70% of RAM. NUM_WORKERS = min(os.cpu_count() or 2, 16) # Batch 128 was measured at ~80% GPU memory on a high-end (80+ GB) GPU — sized for an A100, H100, or # RTX PRO 6000. On a smaller GPU replace this with "auto" and the probe finds the largest batch that fits. BATCH_SIZE = 128 model = RFDETRNano() # type: ignore[no-untyped-call]

5 - Train¶

model.train() builds the TrainConfig and hands the rest to PyTorch Lightning: forward pass, bipartite matching loss, weight updates, learning-rate scheduling, and periodic COCO mAP validation.

CSVLogger appends one row of metrics per epoch to output_dir/metrics.csv, plotted below. The progress bar reports elapsed time per epoch — use the first completed epoch to size how long the remaining 39 will take on your GPU before deciding whether to let this run uninterrupted or resume it across multiple sessions.

To resume an interrupted run, add resume=f"{OUTPUT_DIR}/last.ckpt" to the call below and re-run this cell. last.ckpt is rewritten every epoch (the archive checkpoint_<epoch>.ckpt files follow checkpoint_interval, 10 by default), so a disconnect costs at most the epoch in flight.

The five non-default arguments¶

• batch_size=128 — measured at ~80% of GPU memory on an 80+ GB GPU, leaving headroom for the multi-scale ladder's largest resolution and for memory fragmentation over a long run. Two things to know when changing it: the default lr=1e-4 was tuned around an effective batch of 16, and this notebook leans on cosine-plus-warmup rather than rescaling lr for the 8x larger batch — if you tune, the linear-scaling heuristic is the place to start. And on a smaller GPU, set batch_size="auto": the probe finds the largest safe micro-batch and raises grad_accum_steps toward an effective batch of 16 ("auto" works only through model.train(); building the Lightning modules by hand requires a concrete integer).
• num_workers — the default is 2, which is fine for the few-thousand-image datasets in the other cookbooks and far too low for COCO2017: 118k JPEGs have to be decoded and resized every epoch, and two worker processes cannot keep a modern GPU fed. But more is not simply better — on a 48-core host with this exact configuration, 48 workers sat at ~15% CPU while eating ~70% of RAM: once the GPU is the bottleneck, extra workers add nothing but memory (each worker process carries its own copy of the 118k-image annotation index, plus prefetch_factor batches of decoded tensors queued per worker). Capping at 16 feeds the GPU with headroom; if the GPU ever waits on data, raise the cap until epoch time stops improving.
• seed=0 — seeds Python, NumPy, and torch through Lightning's seed_everything(..., workers=True), which also gives each dataloader worker a distinct, reproducible stream. Note that RF-DETR's own seed_all helper is not used here: besides seeding, it sets cudnn.deterministic=True and torch.use_deterministic_algorithms(True), forcing slow deterministic kernels for exactly the scatter and grid-sample backward passes deformable attention leans on. Reproducible seeding is worth having; deterministic kernels are not, in a notebook about throughput.
• lr_scheduler="cosine" and warmup_epochs=1 — the default schedule is "step" with lr_drop=100, i.e. a 10x drop after epoch 100. In a 40-epoch run that drop never happens and the learning rate stays flat from first step to last, so the mAP curve plateaus noisily with no final refinement. Cosine annealing is sized from the run's own total step count, so it decays correctly no matter how many epochs or how large a batch the GPU ends up with, and the one-epoch linear warmup keeps the first steps stable.

Levers deliberately not pulled¶

Two further speedups exist but change the recipe rather than the schedule, so this notebook leaves them alone. With multi_scale=True (the default) and do_random_resize_via_padding=False, training runs at the largest scale of the multi-scale ladder — 544 px for Nano, while validation and inference stay at 384 px. Setting multi_scale=False trains at 384 px instead (roughly half the pixels) and additionally unblocks RFDETRNano(compile=True), which is skipped whenever multi-scale is on because the varying input shapes retrigger compilation. Separately, augmentation_backend="gpu" moves augmentation off the CPU workers — a rescue for CPU-starved runtimes (a 2-core Colab), and exactly backwards on this notebook's target: with the CPU mostly idle and the GPU saturated, it would add work to the bottleneck and take it away from idle cores.

In [ ]:

model.train(
    dataset_file="coco",
    dataset_dir=COCO_ROOT,
    output_dir=OUTPUT_DIR,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    seed=0,
    lr_scheduler="cosine",
    warmup_epochs=1,
    tensorboard=False,
    progress_bar="tqdm",
)

model.train( dataset_file="coco", dataset_dir=COCO_ROOT, output_dir=OUTPUT_DIR, epochs=EPOCHS, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, seed=0, lr_scheduler="cosine", warmup_epochs=1, tensorboard=False, progress_bar="tqdm", )

6 - Plot CSVLogger metrics¶

RF-DETR's CSVLogger writes val/mAP_50_95, val/mAP_50, and val/mAP_75 to metrics.csv (plus the train/loss family and ema_ variants), and the plotting helpers below discover those columns automatically. val/mAP_50_95 is the primary metric — standard COCO bounding-box mAP averaged across IoU thresholds 0.50-0.95. val/mAP_50 rises fastest and is the clearest early signal of whether training is progressing at all; val/mAP_75 reflects localization precision, not just whether objects are found.

In [ ]:

from IPython.display import display
from matplotlib import pyplot as plt

METRICS_CSV = f"{OUTPUT_DIR}/metrics.csv"

loss_figure = plot_loss_metrics(METRICS_CSV)
display(loss_figure)
plt.close(loss_figure)

from IPython.display import display from matplotlib import pyplot as plt METRICS_CSV = f"{OUTPUT_DIR}/metrics.csv" loss_figure = plot_loss_metrics(METRICS_CSV) display(loss_figure) plt.close(loss_figure)

In [ ]:

map_figure = plot_map_metrics(METRICS_CSV)
display(map_figure)
plt.close(map_figure)

map_figure = plot_map_metrics(METRICS_CSV) display(map_figure) plt.close(map_figure)