How to Load a Trained RF-DETR Model from a Checkpoint: Complete Guide

Load trained RF-DETR models using RFDETRModelModule.load_from_checkpoint(), which automatically handles both Lightning .ckpt files and legacy .pth checkpoints while interpolating positional embeddings to match your model resolution.

The RF-DETR repository by Roboflow provides a PyTorch Lightning-based training framework that simplifies checkpoint management. Whether you are resuming training, running inference, or migrating from older model versions, understanding how to properly load a trained RF-DETR model from a checkpoint ensures you maintain model state, optimizer configurations, and Exponential Moving Average (EMA) weights.

Understanding RF-DETR Checkpoint Architecture

RF-DETR models are wrapped in RFDETRModelModule, a subclass of pytorch_lightning.LightningModule defined in src/rfdetr/training/module_model.py. This wrapper handles the complexity of checkpoint serialization and deserialization, making it straightforward to load a trained RF-DETR model from a checkpoint regardless of the file format.

Checkpoint Contents and State Dicts

When you save or load a checkpoint, the file contains several critical components:

• Model state_dict: The trained weights of the underlying RF-DETR architecture
• Optimizer state: Training optimization parameters (when resuming training)
• EMA state: Exponential Moving Average weights for improved inference performance
• Positional embeddings: DINOv2-derived position encodings that may require interpolation when loading across different image resolutions

The on_load_checkpoint hook in src/rfdetr/training/module_model.py (lines 1448-1505) orchestrates the restoration process, automatically detecting file formats and handling necessary transformations.

Loading Checkpoints with Lightning

The primary method for loading a trained RF-DETR model from a checkpoint uses PyTorch Lightning's standardized API, which RF-DETR extends to handle legacy formats transparently.

Standard Checkpoint Loading

For checkpoints produced by the current RF-DETR training pipeline (.ckpt files), use the class method load_from_checkpoint():

from rfdetr.training.module_model import RFDETRModelModule

# Load the checkpoint

ckpt_path = "path/to/checkpoint.ckpt"
module = RFDETRModelModule.load_from_checkpoint(ckpt_path)

# Extract the underlying model

model = module.model
model.eval()  # Set to evaluation mode for inference

This method invokes the on_load_checkpoint hook at line 1448 of src/rfdetr/training/module_model.py, which restores self.model along with any required positional-embedding interpolation.

Legacy .pth File Support

RF-DETR automatically detects and converts legacy .pth checkpoints produced by earlier training implementations. The conversion happens transparently within on_load_checkpoint (lines 1454-1461):

from rfdetr.training.module_model import RFDETRModelModule

# Load legacy .pth checkpoint - conversion is automatic

legacy_path = "path/to/legacy_checkpoint.pth"
module = RFDETRModelModule.load_from_checkpoint(legacy_path)

model = module.model
model.eval()

During this process, the system interpolates DINOv2 positional embeddings to match the current model resolution, ensuring compatibility even when loading weights trained on different input sizes.

Handling EMA Weights and Special Cases

Exponential Moving Average (EMA) weights often provide superior inference performance compared to raw training weights. RF-DETR's checkpoint system preserves these weights and restores them when appropriate callbacks are configured.

Restoring Exponential Moving Average Weights

When a converted legacy checkpoint contains an EMA state dict, on_load_checkpoint stashes it under legacy_ema_state_dict (lines 1497-1505 of src/rfdetr/training/module_model.py). The RFDETREMACallback class in src/rfdetr/training/callbacks/ema.py then applies these weights:

from rfdetr.training.module_model import RFDETRModelModule
from rfdetr.training.callbacks.ema import RFDETREMACallback
from pytorch_lightning import Trainer

# Initialize trainer with EMA callback

trainer = Trainer(callbacks=[RFDETREMACallback()])

# Load checkpoint containing EMA weights

module = RFDETRModelModule.load_from_checkpoint("ckpt_with_ema.ckpt")

# EMA weights are automatically applied during trainer setup or when resuming training

trainer.fit(module)

Manual Conversion and CLI Usage

While automatic conversion handles most scenarios, you may need to manually convert legacy checkpoints or load models via the command-line interface.

Converting Legacy Checkpoints Manually

For scenarios requiring explicit format conversion without immediate loading, use the utility function in src/rfdetr/training/checkpoint.py:

from rfdetr.training.checkpoint import convert_legacy_checkpoint

# Convert .pth to Lightning-compatible .ckpt

convert_legacy_checkpoint(
    old_checkpoint_path="legacy.pth",
    new_checkpoint_path="converted.ckpt"
)

This produces a PTL-compatible checkpoint that can be loaded by any standard PyTorch Lightning workflow.

Loading via Command Line Interface

The RF-DETR CLI trainer defined in src/rfdetr/training/trainer.py accepts a --ckpt_path argument for seamless checkpoint resumption:

rfdetr train \
  --ckpt_path path/to/checkpoint.ckpt \
  --config configs/rfdetr_small.yaml

The trainer handles checkpoint loading internally, automatically detecting whether you are resuming training or performing inference.

Positional Embedding Interpolation Details

When loading checkpoints across different image resolutions, RF-DETR interpolates DINOv2 positional embeddings to match the target resolution. This interpolation logic, implemented in src/rfdetr/models/weights.py, ensures that checkpoints trained at one resolution can be fine-tuned or evaluated at another without architectural mismatches.

The on_load_checkpoint method specifically handles this interpolation (lines 1454-1461) when it detects resolution differences between the saved checkpoint and the current model configuration.

Summary

• Use RFDETRModelModule.load_from_checkpoint() as the primary API to load a trained RF-DETR model from a checkpoint, located in src/rfdetr/training/module_model.py.
• Automatic format detection handles both modern .ckpt files and legacy .pth checkpoints without code changes.
• EMA weight restoration requires attaching RFDETREMACallback from src/rfdetr/training/callbacks/ema.py to your trainer.
• Positional embedding interpolation occurs automatically in on_load_checkpoint when loading checkpoints trained at different resolutions.
• Manual conversion is available via convert_legacy_checkpoint() in src/rfdetr/training/checkpoint.py for preprocessing legacy files.
• CLI loading uses the --ckpt_path flag with the trainer defined in src/rfdetr/training/trainer.py.

Frequently Asked Questions

What is the difference between .ckpt and .pth files in RF-DETR?

The .ckpt files are PyTorch Lightning checkpoints that contain the full training state including optimizer parameters and metadata, while legacy .pth files contain only model weights. RF-DETR's on_load_checkpoint hook in src/rfdetr/training/module_model.py automatically converts .pth files to the Lightning format during loading, interpolating positional embeddings as needed.

How do I load a checkpoint for inference only?

Load the checkpoint using RFDETRModelModule.load_from_checkpoint(), extract module.model, and call model.eval() to set dropout and batch normalization layers to evaluation mode. This ensures consistent inference behavior without modifying the model weights.

What happens if my checkpoint contains EMA weights?

When loading a checkpoint with EMA weights, RFDETRModelModule stashes the EMA state dict as legacy_ema_state_dict during the on_load_checkpoint hook. If you attach RFDETREMACallback to your Lightning Trainer, it automatically detects and applies these EMA weights, typically providing better inference accuracy than standard training weights.

Can I resume training from a checkpoint saved at a different image resolution?

Yes. The checkpoint loading system in src/rfdetr/models/weights.py automatically interpolates DINOv2 positional embeddings when it detects resolution mismatches. This allows you to fine-tune models on datasets with different image sizes without losing the benefits of pre-trained position encodings.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}