Train RF-DETR Nano on COCO2017 (single GPU)¶

Downloads the full public COCO2017 dataset and trains RF-DETR Nano from scratch for 40 epochs, end to end — pretrain_weights=None, so no released detector checkpoint is loaded (section 4 covers what that costs in accuracy).

This notebook uses torch.compile with automatic mixed precision (bf16 when supported, otherwise fp16) on any CUDA GPU by default. Measured on an RTX PRO 6000 (96 GB) an epoch takes about 7 minutes, so the 40-epoch run fits in about 5 hours; on an L4 (24 GB) an epoch takes about 36 minutes and the run needs several sessions with resume. A100, T4 and other CUDA GPUs run the same recipe with batch_size adjusted; CPU and Apple MPS do not benefit from compilation and should use the fine-tuning cookbooks instead.

An opt-in Transformer Engine FP8 CUDA-graph route is also available (USE_TE_GRAPHS = True in section 4): it nearly doubled throughput in the batch-4 synthetic experiment below, but its full-COCO batch-128 throughput and accuracy still need validation, so it stays off by default. Enabling every optimization together is not supported: FP8 graphs use Transformer Engine capture, not the compiler's graph runtime, so USE_TE_GRAPHS=True disables compile.

Skip the RAM-disk move on a low-RAM host — it needs ~19 GB free in /dev/shm.

In addition to library throughput improvements, this notebook enables compilation, with host-adaptive knobs (num_workers, seed), a batch size measured on two GPUs, 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.
• Measured end to end: the same recipe on rfdetr 1.9.0 takes 65 min/epoch on an RTX PRO 6000 and 97 min on an L4. The 1.9.0 pipeline was not GPU-bound (32 vs 22 img/s on those two cards); develop runs at 289 vs 56 img/s.

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. Measured per-epoch time (training plus validation), for RF-DETR Nano at a fixed 384 px with seed=0, torch 2.11.0+cu128 and Transformer Engine 2.18/2.19, on 11-12 Sep 2026 — RTX PRO 6000 Blackwell 96 GB is the mean of epochs 1-2, L4 24 GB is epoch 1:

GPU	rfdetr 1.9.0 (bf16, batch at ceiling)	develop bf16 + compile	develop FP8 + compile
RTX PRO 6000	65 min @ 128 (160 OOM)	7.2 min @ 128 (40 % mem), 7.3 min @ 288	7.0 min @ 128, 7.1 min @ 320
L4	97 min @ 32 (40 is slower)	35 min @ 32 (46 % mem), 36 min @ 64	40 min @ 32, 37 min @ 72

Epoch 0 is slower (compile warm-up, annotation index, page cache): about 8 % on the RTX PRO 6000 and up to 40 % on an L4. Time epoch 1 onward.

Session limits vary; do not assume all 40 epochs will fit in one session. Training still checkpoints every checkpoint_interval epochs (10 by default) to last.ckpt as a safety net: if a session merely disconnects and the same runtime is still alive, 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. But OUTPUT_DIR below is local runtime disk: once Colab recycles the runtime (idle timeout, 12/24h cap, or a paid-tier disconnect), last.ckpt is gone with it, and there is nothing left to resume from. For a run that genuinely needs several sessions, point OUTPUT_DIR at mounted Google Drive (from google.colab import drive; drive.mount("/content/drive"), then OUTPUT_DIR = "/content/drive/MyDrive/rfdetr_coco2017") instead of local disk, or copy last.ckpt off the runtime before it ends and copy it back before 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 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 the git URL with a version pin. The install below also includes the CUDA/Transformer Engine dependency, needed only if you opt into the FP8 graph route (USE_TE_GRAPHS = True in section 4); this cookbook change and its runtime implementation must land together. Older revisions can warn and run FP8 eagerly instead. If you enable it, confirm the training log reports Backend: Transformer Engine FP8.

GPU required. Training needs a CUDA GPU — in Colab: Runtime → Change runtime type → GPU. The default recipe below is BF16/FP16 compilation; the optional Transformer Engine FP8 graph route is covered in section 4.

In [ ]:

%pip install --progress-bar on --no-build-isolation "rfdetr[train,augment,visual,cuda] @ git+https://github.com/roboflow/rf-detr.git@develop"
%pip install --progress-bar on "jedi>=0.18"
import sys

print("Using interpreter:", sys.executable)
!{sys.executable} -m pip check || echo "warning: pip check reported conflicts (may be pre-existing/unrelated)"

%pip install --progress-bar on --no-build-isolation "rfdetr[train,augment,visual,cuda] @ git+https://github.com/roboflow/rf-detr.git@develop" %pip install --progress-bar on "jedi>=0.18" import sys print("Using interpreter:", sys.executable) !{sys.executable} -m pip check || echo "warning: pip check reported conflicts (may be pre-existing/unrelated)"

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¶

pretrain_weights=None skips the released COCO-pretrained Nano checkpoint, so the detection transformer and its heads start from random initialization — this notebook trains COCO2017 from scratch. The DINOv2 backbone still loads its self-supervised hub weights: RF-DETR requests them precisely when pretrain_weights is None, because a full checkpoint would otherwise carry the backbone with it.

Instantiating with pretrain_weights=None raises a PretrainWeightsCompatibilityWarning saying the model is initialised from scratch. That is expected here; it is a warning, not an error.

Accuracy expectation. From scratch in 40 epochs lands far below the published Nano checkpoint. The released weights come from a much longer schedule with large-scale detection pretraining behind them. Fine-tune from the default checkpoint (drop the pretrain_weights argument) whenever accuracy, not the training loop itself, is the goal.

num_classes is left at its default of 90 — the standard COCO category-id space RF-DETR's dataset_file="coco" loader already uses, 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 -> GPU.")

print("PyTorch:", torch.__version__, "CUDA:", torch.version.cuda)
print("GPU:", torch.cuda.get_device_name(0))

# 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
USE_TE_GRAPHS = False  # Opt-in: True selects the Transformer Engine FP8 CUDA-graph showcase (RTX PRO 6000).
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 saturated the measured RTX PRO 6000 BF16 baseline; FP8 graph memory/throughput needs validation.
# Other GPUs: use the configuration table below.
BATCH_SIZE = 128  # User-tune this for the available GPU; Transformer Engine evidence below uses synthetic batch 4 only.
# FP8 uses Transformer Engine graphs. The compilation alternative enables Inductor graphs only at small batches.
CUDA_GRAPHS = USE_TE_GRAPHS or BATCH_SIZE <= 16

model = RFDETRNano(
    compile=not USE_TE_GRAPHS,
    cuda_graphs=CUDA_GRAPHS,
    pretrain_weights=None,
)  # 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 -> GPU.") print("PyTorch:", torch.__version__, "CUDA:", torch.version.cuda) print("GPU:", torch.cuda.get_device_name(0)) # 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 USE_TE_GRAPHS = False # Opt-in: True selects the Transformer Engine FP8 CUDA-graph showcase (RTX PRO 6000). 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 saturated the measured RTX PRO 6000 BF16 baseline; FP8 graph memory/throughput needs validation. # Other GPUs: use the configuration table below. BATCH_SIZE = 128 # User-tune this for the available GPU; Transformer Engine evidence below uses synthetic batch 4 only. # FP8 uses Transformer Engine graphs. The compilation alternative enables Inductor graphs only at small batches. CUDA_GRAPHS = USE_TE_GRAPHS or BATCH_SIZE <= 16 model = RFDETRNano( compile=not USE_TE_GRAPHS, cuda_graphs=CUDA_GRAPHS, pretrain_weights=None, ) # type: ignore[no-untyped-call]

GPU configuration recommendations¶

These settings target throughput for Nano at fixed 384 px, not optimal accuracy. Set USE_TE_GRAPHS and BATCH_SIZE in the preceding cell before constructing the model; the precision, compile, and graph flags follow automatically. Starting point means not benchmarked on that GPU; it is not a fastest-path claim.

GPU	BATCH_SIZE	USE_TE_GRAPHS	Precision	Compile / graphs	Evidence and recommendation
RTX PRO 6000 Blackwell, 96 GB	128	False	BF16	on / off	Default; measured large-batch COCO baseline.
RTX PRO 6000 Blackwell, opt-in FP8	128	True	FP8	off / on	Opt-in showcase; graph gain measured only at synthetic batch 4. Validate full-COCO throughput and accuracy before adopting.
L4, 24 GB	32	False	BF16	on / off	Measured throughput saturates around batch 32; ordinary FP8 was slower. FP8 graphs at this batch are unmeasured.
H100, 80 GB	128	False	BF16	on / off	Starting point; FP8 is supported, but this workload has no matched FP8-graph benchmark here.
A100, 40/80 GB	64	False	BF16	on / off	Starting point; Ampere does not support this FP8 route.
RTX 6000 Ada, 48 GB	64	False	BF16	on / off	Starting point; FP8 capable, but not the Blackwell GPU used in the uploaded graph experiment.
RTX A6000 / A10	32	False	BF16	on / off	Conservative starting point; Ampere, no FP8 route.
T4 / Quadro RTX 6000	32	False	FP16	on / off	Starting point; older Turing GPUs, no BF16/FP8 route. Reduce batch if memory is insufficient.

GPU identity comes from NVIDIA's architecture table; Transformer Engine supports FP8 on Ada, Hopper and Blackwell. All rows retain 40 epochs, grad_accum_steps=1, multi_scale=False, square_resize_div_64=True, and num_workers=min(os.cpu_count() or 2, 16). On CPU-constrained hosts, worker count and JPEG decoding can dominate GPU differences. Measure complete epochs after startup, memory, and validation mAP; choose the fastest configuration meeting the same accuracy target. Lowering batch size changes effective batch size.

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. Compilation adds startup overhead: time subsequent epochs separately. A failed CUDA graph capture (CUDA_GRAPHS=True or USE_TE_GRAPHS=True) requires a runtime restart before retrying. The Transformer Engine graph path, when enabled, captures one fixed batch/resolution signature; changing that shape raises.

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.

Training settings¶

• batch_size=128 — throughput stops improving once the GPU is saturated: on the RTX PRO 6000, 128 and 288 give the same images per second; on the L4, 32 and 64 do. Pick the batch that leaves memory headroom rather than the largest that fits. 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. On a smaller GPU, reduce the integer batch size and adjust grad_accum_steps to preserve the desired effective batch; in bf16, batch_size="auto" also works and probes for the largest micro-batch that fits.
• 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.

Compilation¶

RFDETRNano(compile=True) is worth 21 % per epoch on the RTX PRO 6000 and 27 % on an L4 at the same batch, precision and resolution. Epoch 0 pays the warm-up. Inductor's coalesce tiling analysis is disabled by the library at the compile call (torch 2.11 asserts on it and would silently fall back to eager for the whole forward), so no user-side flag is needed.

multi_scale=False trains Nano at a fixed 384 px rather than the multi-scale ladder's 544 px top scale, which is the resolution every number above was measured at; it replaces the previous recipe's 544 px workload and may change accuracy.

CUDA graphs on top of compilation¶

RFDETRNano(compile=True, cuda_graphs=True) replays the compiled forward and backward through Inductor's CUDA graph trees, removing the kernel-launch gaps that compilation leaves in place. Whether that is worth anything depends on the batch, measured on the RTX PRO 6000 with Nano at 384 px, BF16, synthetic batches:

batch [img]	eager [img/s]	compile [img/s]	compile + cuda_graphs [img/s]	GPU busy compile → both [-]
4	78.7	96.8	116.0 (1.20x over compile)	0.32 → 0.42
64	246.8	322.8	325.7 (+0.9 %, noise)	0.82 → 0.83

With USE_TE_GRAPHS=False (the default) and batch 128, the compiled kernels already run back to back, so this recipe leaves Inductor graphs off. It enables them at batches up to 16, where launch gaps can matter. This rule does not affect the optional Transformer Engine graphs, controlled separately by USE_TE_GRAPHS. The compiled graph route is validated for single-GPU detection without grad_accum_steps > 1; other layouts fall back to plain compilation with a warning. Graph recording adds a few seconds per new input shape on top of compile warm-up. See Combining CUDA graphs with compilation.

What is left¶

With compilation on, the GPU is busy for about 83 % of each training step at batch 64 and the remaining share is spent outside the model: bipartite matching and the loss on the host, and the data loader. No constructor or train() flag in this notebook moves that share — larger batches, more workers, standard FP8, and Inductor CUDA graphs were each measured and none reach it. The optional Transformer Engine graph path has only synthetic fixed-shape evidence, not an end-to-end COCO measurement. The one wall-clock knob that remains is validation: eval_interval=N runs the 5,000- image COCO evaluation every N epochs instead of every epoch, at the cost of a coarser mAP curve in section 6; its saving was not measured for this recipe.

Optional: Transformer Engine FP8 graphs¶

Requires the cuda extra (Transformer Engine, Hopper/Ada/Blackwell, Linux x86-64), an integer batch_size, and amp_dtype="fp8". Measured on Nano, FP8 changes throughput by +3 % per epoch on Blackwell and −11 % on an L4, and raises the memory ceiling to 320 from 288 (RTX PRO 6000) and to 72 from 64 (L4). Lightning warns at startup that 14 linear layers are not FP8-shaped; those are the per-layer classification heads and they run in bf16. Those ordinary-FP8 measurements exclude Transformer Engine graphs; they do not predict this combined route. The setup cell above installs transformer-engine[pytorch] through RF-DETR's cuda extra and builds its PyTorch extension: CUDA toolkit headers, cuDNN, and a compiler have to be present, Blackwell needs CUDA 12.8 or later, and the kernel has to be restarted afterwards. See the advanced FP8 setup guide.

The install cell above uses RF-DETR's existing cuda extra for Transformer Engine and runs pip check after installing Jedi. The Transformer Engine-aware graph path was tested with Transformer Engine 2.19.0; no separate CUDA-core package or version override is required. Transformer Engine must match the CUDA toolkit, driver, cuDNN, and PyTorch ABI; package resolution alone is not compatibility evidence. Restart the kernel after installing compiled dependencies, set USE_TE_GRAPHS = True, and run the full 40-epoch training cell with amp_dtype="fp8" below.

This path is limited to single-GPU detection with compile=False, grad_accum_steps=1, multi_scale=False and square_resize_div_64=True; aspect-preserving batches with the last setting disabled stay eager. It captures one fixed batch/resolution signature; unsupported layouts warn and stay eager. A synthetic Nano 384 px / batch-4 run on an RTX PRO 6000 Blackwell used 20 warm-up and 50 measured steps: 78.075 ms eager versus 40.066 ms with Transformer Engine-aware graphs (about 1.95x throughput), with one capture serving 70 calls. This is not COCO parity, accuracy, or large-batch evidence; repeat numerical-parity and performance checks on the target GPU before adopting it as a baseline.

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",
    amp_dtype="fp8" if USE_TE_GRAPHS else "auto",
    grad_accum_steps=1,
    multi_scale=False,
    square_resize_div_64=True,
)

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", amp_dtype="fp8" if USE_TE_GRAPHS else "auto", grad_accum_steps=1, multi_scale=False, square_resize_div_64=True, )

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)