Feature request: infer model size from checkpoint for easier loading

[pirnerjonas]

Search before asking

• I have searched the RF-DETR issues and found no similar feature requests.

Description

Problem

When loading a fine-tuned RF-DETR checkpoint, the user must know and specify the correct model size class (RFDETRLarge, RFDETRBase, etc.). This becomes inconvenient when you have old checkpoints on disk and don't remember which variant was used for training.

Currently you have to do:

from rfdetr import RFDETRLarge  # must know this upfront
model = RFDETRLarge(pretrain_weights="checkpoint_best_total.pth", num_classes=2)

Observation

The checkpoint already stores this information. The saved args namespace contains pretrain_weights which includes the model size in the filename (e.g. rf-detr-large-2026.pth), and encoder which differs by variant. For example:

import torch
ckpt = torch.load("checkpoint_best_total.pth", map_location="cpu", weights_only=False)
print(ckpt["args"].pretrain_weights)  # 'rf-detr-large-2026.pth'
print(ckpt["args"].encoder)           # 'dinov2_windowed_small'

Suggestion

It would be nice to have a convenience loader that infers the model class from the checkpoint, e.g.:

import rfdetr
model = rfdetr.from_checkpoint("checkpoint_best_total.pth")

This could read args from the checkpoint, determine the correct model class, and instantiate it — without the user needing to remember or specify the variant.

Alternatively, storing an explicit model_size or model_class field in args at training time would make this even more robust than relying on filename parsing.

Use case

No response

Additional

No response

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[ashwin3005]

Facing the same issue,

currently I am loading the checkpoints (.pth) using torch.load(), then getting the model details like "number of decoder layers", "hidden_dim" etc.. from state_dict. and comparing that with pre-filled table for choosing the class for appropriate variant.

[ashwin3005]

Facing the same issue,

currently I am loading the checkpoints (.pth) using torch.load(), then getting the model details like "number of decoder layers", "hidden_dim" etc.. from state_dict. and comparing that with pre-filled table for choosing the class for appropriate variant.

this is how I am doing it.

import torch

# Model specifications (based on RF-DETR 1.4.0+)

MODEL_SPECS = {
    "nano": (2, 256, 16, 577),
    "small": (3, 256, 16, 1025),
    "medium": (4, 256, 16, 1297),
    "base": (3, 256, 14, 1370),
    "large": (4, 256, 16, 1937),
    "xlarge": (5, 512, 20, 1226),
    "xxlarge": (5, 512, 20, 1937),
}  

# the above dict is hard-coded after inspecting the model architecture of every model

weights_path = "checkpoint_best_total.pth"

# Load the model state
state = torch.load(weights_path, map_location="cpu", weights_only=False)
model_sd = state["model"]

# extract the signature from the model state (decoder_layers, hidden_dim, patch_size, pe_size) 

decoder_layers = {int(k.split(".")[3]) for k in model_sd if k.startswith("transformer.decoder.layers.")}
num_decoder_layers = max(decoder_layers) + 1

hidden_dim = model_sd["transformer.decoder.layers.0.self_attn.in_proj_weight"].shape[1]

patch_size = model_sd["backbone.0.encoder.encoder.embeddings.patch_embeddings.projection.weight"].shape[-1]

pe_size = model_sd["backbone.0.encoder.encoder.embeddings.position_embeddings"].shape[1]

signature = (num_decoder_layers, hidden_dim, patch_size, pe_size)

# Find the corresponding RF-DETR variant
variant = None
for name, spec in MODEL_SPECS.items():
    if spec == signature:
        variant = name
        break

if variant:
    print(f"Detected RF-DETR variant: {variant}")
else:
    raise RuntimeError("Unknown RF-DETR variant detected.")

do let me know, if there is a better way to do this.