Skip to content

RF-DETR Medium

Bases: RFDETR

Train an RF-DETR Medium model.

Training accepts custom square integer resolution values. The value must be divisible by patch_size * num_windows.

Source code in src/rfdetr/variants.py
class RFDETRMedium(RFDETR):
    """Train an RF-DETR Medium model.

    Training accepts custom square integer ``resolution`` values. The value must be divisible by ``patch_size *
    num_windows``.
    """

    size = "rfdetr-medium"
    _model_config_class = RFDETRMediumConfig

Attributes

class_names property

Retrieve the class names supported by the loaded model.

Returns:

Type Description
list[str]

A list of class name strings, 0-indexed. When no custom class names are embedded in the checkpoint, returns

list[str]

the standard 80 COCO class names.

is_optimized_inplace property

Report destructive optimization state retained by a legacy RFDETR instance.

The snapshot factory :meth:inference never sets this flag on its source.

Functions

__init__(*, trust_checkpoint=False, **kwargs)

Initialize with ModelConfig fields as keyword arguments.

Passes all remaining kwargs to the variant's ModelConfig. Unknown kwargs raise pydantic.ValidationError. See the variant's config class for available parameters (e.g. RFDETRSmallConfig).

Parameters:

Name Type Description Default

trust_checkpoint

bool

When True, allow pretrain_weights to fall back to full pickle deserialization (weights_only=False) if safe loading fails. Set this only when pretrain_weights points to a checkpoint you explicitly trust (e.g. when called internally by :meth:from_checkpoint); left False by default for the ordinary construction path, which only ever downloads official Roboflow-hosted weights.

False

**kwargs

Any

ModelConfig field values (e.g. resolution, num_classes, pretrain_weights, gradient_checkpointing).

{}

Raises:

Type Description
ImportError

If backbone_lora=True is set, or pretrain_weights was saved by a backbone_lora=True run, and peft is not installed (pip install "rfdetr[lora]").

Source code in src/rfdetr/detr.py
def __init__(self, *, trust_checkpoint: bool = False, **kwargs: Any) -> None:
    """Initialize with ModelConfig fields as keyword arguments.

    Passes all remaining kwargs to the variant's ModelConfig. Unknown kwargs raise
    ``pydantic.ValidationError``. See the variant's config class for available
    parameters (e.g. ``RFDETRSmallConfig``).

    Args:
        trust_checkpoint: When ``True``, allow ``pretrain_weights`` to fall back to full
            pickle deserialization (``weights_only=False``) if safe loading fails. Set this
            only when ``pretrain_weights`` points to a checkpoint you explicitly trust (e.g.
            when called internally by :meth:`from_checkpoint`); left ``False`` by default for
            the ordinary construction path, which only ever downloads official Roboflow-hosted
            weights.
        **kwargs: ModelConfig field values (e.g. ``resolution``, ``num_classes``,
            ``pretrain_weights``, ``gradient_checkpointing``).

    Raises:
        ImportError: If ``backbone_lora=True`` is set, or ``pretrain_weights`` was saved by a
            ``backbone_lora=True`` run, and ``peft`` is not installed (``pip install "rfdetr[lora]"``).
    """
    self.model_config = self.get_model_config(**kwargs)
    self.maybe_download_pretrain_weights()
    self.model = self.get_model(self.model_config, trust_checkpoint=trust_checkpoint)
    self.callbacks: dict[str, list[Callable[..., Any]]] = defaultdict(list)

    self.means = list(self.means)
    self.stds = list(self.stds)

    # repeat means and stds for non-rgb images
    if self.model_config.num_channels != 3:
        from itertools import cycle

        self.means = [val for _, val in zip(range(self.model_config.num_channels), cycle(self.means))]
        self.stds = [val for _, val in zip(range(self.model_config.num_channels), cycle(self.stds))]

    self.model.inference_model = None
    self._is_optimized_for_inference = False
    self._has_warned_about_not_being_optimized_for_inference = False
    self._optimized_has_been_compiled = False
    self._optimized_batch_size: int | None = None
    self._optimized_resolution: int | None = None
    self._optimized_dtype: torch.dtype | None = None
    self._optimized_inplace = False
    self._has_been_trained = False

deploy_to_roboflow(workspace, project_id, version=None, api_key=None, size=None)

Deploy the trained RF-DETR model to Roboflow.

Deploying with Roboflow will create a Serverless API to which you can make requests.

You can also download weights into a Roboflow Inference deployment for use in Roboflow Workflows and on-device deployment.

Parameters:

Name Type Description Default

workspace

str

The name of the Roboflow workspace to deploy to.

required

project_id

str

The project ID to which the model will be deployed.

required

version

int | str | None

The project version to which the model will be deployed. If not provided, the highest existing dataset version is resolved automatically via the Roboflow API; for a project with no generated versions yet the lookup falls back to version 1, and the Roboflow SDK then raises its own RuntimeError ("Version number 1 is not found.").

None

api_key

str | None

Your Roboflow API key. If not provided, it will be read from the environment variable ROBOFLOW_API_KEY.

None

size

str | None

The size of the model to deploy. If not provided, it will default to the size of the model being trained (e.g., "rfdetr-base", "rfdetr-large", etc.).

None

Raises:

Type Description
ValueError

If the api_key is not provided and not found in the environment variable ROBOFLOW_API_KEY, or if the size is not set for custom architectures.

RuntimeError

If a legacy in-place optimization cleared the base weights.

Note

Bundle creation is delegated to :meth:export_for_roboflow, which can be called independently to write weights.pt and class_names.txt without a network round-trip.

Source code in src/rfdetr/detr.py
def deploy_to_roboflow(
    self,
    workspace: str,
    project_id: str,
    version: int | str | None = None,
    api_key: str | None = None,
    size: str | None = None,
) -> None:
    """Deploy the trained RF-DETR model to Roboflow.

    Deploying with Roboflow will create a Serverless API to which you can make requests.

    You can also download weights into a Roboflow Inference deployment for use in Roboflow Workflows and on-device
    deployment.

    Args:
        workspace: The name of the Roboflow workspace to deploy to.
        project_id: The project ID to which the model will be deployed.
        version: The project version to which the model will be deployed. If not provided, the highest
            existing dataset version is resolved automatically via the Roboflow API; for a project with
            no generated versions yet the lookup falls back to version ``1``, and the Roboflow SDK then
            raises its own ``RuntimeError`` ("Version number 1 is not found.").
        api_key: Your Roboflow API key. If not provided,
            it will be read from the environment variable `ROBOFLOW_API_KEY`.
        size: The size of the model to deploy. If not provided,
            it will default to the size of the model being trained (e.g., "rfdetr-base", "rfdetr-large", etc.).

    Raises:
        ValueError: If the `api_key` is not provided and not found in the
            environment variable `ROBOFLOW_API_KEY`, or if the `size` is not set for custom architectures.
        RuntimeError: If a legacy in-place optimization cleared the base weights.

    Note:
        Bundle creation is delegated to :meth:`export_for_roboflow`, which can be called independently
        to write ``weights.pt`` and ``class_names.txt`` without a network round-trip.
    """
    if getattr(self, "_optimized_inplace", False) or self.model.model is None:
        raise RuntimeError(
            "Cannot deploy after legacy inplace optimization — "
            "the model weights have been cleared from memory. "
            "Call export_for_roboflow() before optimizing, then deploy the exported bundle."
        )

    from roboflow import Roboflow

    if api_key is None:
        api_key = os.getenv("ROBOFLOW_API_KEY")
        if api_key is None:
            raise ValueError("Set api_key=<KEY> in deploy_to_roboflow or export ROBOFLOW_API_KEY=<KEY>")

    rf = Roboflow(api_key=api_key)
    rf_workspace = rf.workspace(workspace)

    if self.size is None and size is None:
        raise ValueError("Must set size for custom architectures")

    if size is not None and self.size is not None and size != self.size:
        warnings.warn(
            f"deploy_to_roboflow(size={size!r}) overrides this model's own size {self.size!r}; "
            f"deploying as {size!r}. Omit size to deploy with the model's own size.",
            UserWarning,
            stacklevel=2,
        )
    # Explicit user argument wins; fall back to the trained model's size (documented behaviour).
    size = self.size if size is None else size
    with tempfile.TemporaryDirectory(prefix="roboflow_upload_") as tmp_out_dir:
        self.export_for_roboflow(tmp_out_dir)
        project = rf_workspace.project(project_id)
        if version is None:
            # Version ids come back as "<workspace>/<project>/<number>"; the highest number is the
            # newest dataset version. default=1 keeps the SDK's own "Version number 1 is not found."
            # as the error surface for a project with no generated versions.
            version = max(
                (int(os.path.basename(info["id"])) for info in project.get_version_information()),
                default=1,
            )
            logger.info(f"deploy_to_roboflow: no version given, resolved latest version {version}")
        project_version = project.version(version)
        project_version.deploy(model_type=size, model_path=tmp_out_dir, filename="weights.pt")

evaluate(*, split='test', **kwargs)

Evaluate the current model on a dataset split and return COCO metrics.

Runs a single evaluation pass over the requested split via the PyTorch Lightning stack and returns the COCO metrics (mAP, mAR, and the macro-F1 sweep) computed by :class:~rfdetr.training.callbacks.coco_eval.COCOEvalCallback. The same metrics are also printed to the terminal. This works both directly after :meth:train and on a model loaded via :meth:from_checkpoint — the weights already held in memory are evaluated; no checkpoint file is re-loaded.

Apart from split, this method accepts exactly the same keyword arguments as :meth:train (dataset_dir, device, resolution, batch_size, output_dir, num_workers, ...); they are handled identically via the shared :func:_prepare_run_config. This parity is for convenience — the same kwargs dict used for :meth:train can be reused here — not a guarantee every field has an effect. Training-only fields (epochs, lr, weight_decay, ema, early_stopping, run, project, checkpoint_interval, tensorboard/wandb/mlflow/clearml, and similar) are silently accepted and ignored: evaluate() runs through an eval-only trainer (include_training_callbacks=False) that never builds EMA, drop-path, checkpointing, early-stopping, or logger callbacks, so those fields have nothing to attach to.

Unlike :meth:train, this method never adapts the detection head to the dataset: the model is evaluated exactly as configured. If the dataset's class count differs from the model's num_classes a :class:UserWarning is emitted and evaluation proceeds with the model's head unchanged.

Unlike :meth:train, a resolution override does not persist: :attr:model_config (and any cached model.resolution / model.args inference context) is restored to its pre-call values once the eval-only config copy has captured the override, so a later :meth:predict / :meth:export / :meth:train call is unaffected by an evaluate(resolution=...) call.

Parameters:

Name Type Description Default

split

Literal['test', 'val']

Which split to evaluate. "test" evaluates the test/ folder (YOLO-format datasets — plain YOLO datasets and Roboflow exports whose detected format is YOLO — fall back to the valid/ split, with a logged warning, when no test split can be resolved. A Roboflow export in COCO format has no such fallback and still raises FileNotFoundError when its test split is missing; COCO/Objects365 never attempt test and always evaluate valid/) via trainer.test; "val" evaluates the valid/ folder via trainer.validate.

'test'

**kwargs

Any

The same keyword arguments accepted by :meth:train — dataset_dir is required (here or already on the config), and the rest are forwarded to :func:_prepare_run_config / :meth:get_train_config.

{}

Returns:

Type Description
dict[str, float]

Mapping of metric name to value for the evaluated split, e.g. ``{"test/mAP_50_95": ..., "test/mAP_50": ...,

dict[str, float]

"test/F1": ...}. Per-class keys ("test/AP/"``) are included only when

dict[str, float]

log_per_class_metrics=True (default False). Empty when the trainer returns no metrics.

Raises:

Type Description
ImportError

If training dependencies are not installed. Install with pip install "rfdetr[train,loggers]".

ValueError

If split is not "test" or "val".

Source code in src/rfdetr/detr.py
def evaluate(self, *, split: Literal["test", "val"] = "test", **kwargs: Any) -> dict[str, float]:
    """Evaluate the current model on a dataset split and return COCO metrics.

    Runs a single evaluation pass over the requested split via the PyTorch Lightning stack and returns the COCO
    metrics (mAP, mAR, and the macro-F1 sweep) computed by
    :class:`~rfdetr.training.callbacks.coco_eval.COCOEvalCallback`. The same metrics are also printed to the
    terminal. This works both directly after :meth:`train` and on a model loaded via :meth:`from_checkpoint` — the
    weights already held in memory are evaluated; no checkpoint file is re-loaded.

    Apart from ``split``, this method accepts exactly the same keyword arguments as :meth:`train` (``dataset_dir``,
    ``device``, ``resolution``, ``batch_size``, ``output_dir``, ``num_workers``, ...); they are handled identically
    via the shared :func:`_prepare_run_config`. This parity is for convenience — the same kwargs dict used for
    :meth:`train` can be reused here — not a guarantee every field has an effect. Training-only fields (``epochs``,
    ``lr``, ``weight_decay``, ``ema``, ``early_stopping``, ``run``, ``project``, ``checkpoint_interval``,
    ``tensorboard``/``wandb``/``mlflow``/``clearml``, and similar) are silently accepted and ignored: ``evaluate()``
    runs through an eval-only trainer (``include_training_callbacks=False``) that never builds EMA, drop-path,
    checkpointing, early-stopping, or logger callbacks, so those fields have nothing to attach to.

    Unlike :meth:`train`, this method never adapts the detection head to the dataset: the model is evaluated exactly
    as configured. If the dataset's class count differs from the model's ``num_classes`` a :class:`UserWarning` is
    emitted and evaluation proceeds with the model's head unchanged.

    Unlike :meth:`train`, a ``resolution`` override does **not** persist: :attr:`model_config` (and any cached
    ``model.resolution`` / ``model.args`` inference context) is restored to its pre-call values once the
    eval-only config copy has captured the override, so a later :meth:`predict` / :meth:`export` / :meth:`train`
    call is unaffected by an ``evaluate(resolution=...)`` call.

    Args:
        split: Which split to evaluate. ``"test"`` evaluates the ``test/`` folder (YOLO-format datasets — plain
            YOLO datasets and Roboflow exports whose detected format is YOLO — fall back to the ``valid/`` split,
            with a logged warning, when no ``test`` split can be resolved. A Roboflow export in COCO format has
            no such fallback and still raises ``FileNotFoundError`` when its ``test`` split is missing;
            COCO/Objects365 never attempt ``test`` and always evaluate ``valid/``) via ``trainer.test``;
            ``"val"`` evaluates the ``valid/`` folder via ``trainer.validate``.
        **kwargs: The same keyword arguments accepted by :meth:`train` — ``dataset_dir`` is required (here or
            already on the config), and the rest are forwarded to :func:`_prepare_run_config` /
            :meth:`get_train_config`.

    Returns:
        Mapping of metric name to value for the evaluated split, e.g. ``{"test/mAP_50_95": ..., "test/mAP_50": ...,
        "test/F1": ...}``. Per-class keys (``"test/AP/<class>"``) are included only when
        ``log_per_class_metrics=True`` (default ``False``). Empty when the trainer returns no metrics.

    Raises:
        ImportError: If training dependencies are not installed. Install with
            ``pip install "rfdetr[train,loggers]"``.
        ValueError: If ``split`` is not ``"test"`` or ``"val"``.
    """
    from rfdetr.models.weights import interpolate_position_embeddings

    # Training extras (pytorch_lightning et al.) are optional; mirror train()'s import guard.
    try:
        from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
    except ModuleNotFoundError as exc:
        if exc.name and exc.name.startswith("rfdetr."):
            raise
        raise ImportError(
            "RF-DETR training dependencies are missing. "
            'Install them with `pip install "rfdetr[train,loggers]"` and try again.',
        ) from exc

    if split not in ("test", "val"):
        raise ValueError(f"split must be 'test' or 'val', got {split!r}.")

    # Same kwarg handling as train() (device, resolution, deprecated knobs, auto-batch).  A `resolution`
    # override mutates `model_config`/`model.args` in place inside `_prepare_run_config` — intentional and
    # persistent for train(), but evaluate() is documented as inspection-only, so the mutation is snapshotted
    # here and restored once the eval-only config copy below has captured the overridden values it needs.
    _orig_resolution = self.model_config.resolution
    _orig_pe = self.model_config.positional_encoding_size
    _live_model = getattr(self, "model", None)
    _live_args = getattr(_live_model, "args", None) if _live_model is not None else None
    _orig_model_resolution = getattr(_live_model, "resolution", None) if _live_model is not None else None
    _orig_args_resolution = getattr(_live_args, "resolution", None) if _live_args is not None else None
    _orig_args_pe = getattr(_live_args, "positional_encoding_size", None) if _live_args is not None else None
    try:
        config, _accelerator, _devices = _prepare_run_config(self, for_eval=True, **kwargs)

        # Build the module without re-loading pretrain weights, then transplant the
        # already-loaded in-memory weights into it (the reverse of train()'s final
        # `self.model.model = module.model`).  This evaluates the current weights and
        # never passes ``ckpt_path``, sidestepping PTL's loop-state restore on a bare .pth.
        eval_model_config = self.model_config.model_copy(update={"pretrain_weights": None})
    finally:
        self.model_config.resolution = _orig_resolution
        self.model_config.positional_encoding_size = _orig_pe
        if _live_model is not None:
            if hasattr(_live_model, "resolution"):
                _live_model.resolution = _orig_model_resolution
            if _live_args is not None:
                if hasattr(_live_args, "resolution"):
                    _live_args.resolution = _orig_args_resolution
                if hasattr(_live_args, "positional_encoding_size"):
                    _live_args.positional_encoding_size = _orig_args_pe
    # Every weight is transplanted from the live model below, so skip fetching upstream encoder weights.
    module = RFDETRModelModule(eval_model_config, config, load_encoder_weights=False)

    # Free the original model's accelerator memory for the transplant -- otherwise the resident
    # original and the freshly built (randomly initialized) eval module are both on the accelerator
    # simultaneously (peak ~2x model memory), risking OOM on the largest variants right after train()
    # just fit them.  state_dict() below returns references into the live parameter tensors, so this
    # move is reflected in `source_state` regardless of call order; restored immediately once the eval
    # module has consumed the weights, before the (possibly slow) datamodule/trainer setup below.
    _original_device = getattr(self.model, "device", None)
    _moved_to_cpu = _original_device is not None and str(_original_device) != "cpu"
    source_model = self.model.model
    if source_model is None:
        raise RuntimeError("Cannot evaluate: the base model has been cleared by a previous inplace optimization.")
    if _moved_to_cpu:
        # Through ``_locked_move``: this is the live module, so the move must not overlap the deferred
        # first-use move a concurrent ``predict()`` may be running on the same parameters.
        source_model = _locked_move(source_model, "cpu")
        self.model.model = source_model
    try:
        source_state = source_model.state_dict()
        # Reconcile DINOv2 positional embeddings when a `resolution` override changed the PE grid
        # (no-op when unchanged), so the transplant works at the evaluation resolution.
        interpolate_position_embeddings(source_state, eval_model_config.positional_encoding_size)
        module.model.load_state_dict(source_state)
    finally:
        if _moved_to_cpu:
            _move_model_context_to_device(self.model)
    # `self.model_config` was already restored to its pre-call values by the `finally` block above; a
    # `resolution` override only survives on `eval_model_config`, which is what actually built `module`.
    # The datamodule must use the same config so the dataloader resizes to the resolution being evaluated.
    datamodule = RFDETRDataModule(eval_model_config, config)

    # Warn (do not adapt) when the dataset class count differs from the model's head.
    stage = "test" if split == "test" else "validate"
    datamodule.setup(stage)
    dataset_class_names = getattr(datamodule, "class_names", None)
    if isinstance(dataset_class_names, list) and len(dataset_class_names) != self.model_config.num_classes:
        warnings.warn(
            f"Dataset '{config.dataset_dir}' has {len(dataset_class_names)} classes but the model has "
            f"num_classes={self.model_config.num_classes}. Evaluating with the model's head unchanged; "
            "class indices may not line up with the dataset.",
            UserWarning,
            stacklevel=2,
        )

    trainer_kwargs: dict[str, Any] = {"include_training_callbacks": False}
    if _accelerator is not None:
        trainer_kwargs["accelerator"] = _accelerator
    if _devices is not None:
        trainer_kwargs["devices"] = _devices
    trainer = build_trainer(config, self.model_config, **trainer_kwargs)

    # The eval trainer intentionally has no logger; the metric callback still logs with
    # ``logger=True``, so PTL emits one "no logger configured" warning per metric.  Suppress
    # only that specific message — the metrics are still collected from the trainer's results.
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*no logger configured.*")
        if split == "test":
            results = trainer.test(module, datamodule)
        else:
            results = trainer.validate(module, datamodule)

    return {key: float(value) for key, value in results[0].items()} if results else {}

export(output_dir='output', infer_dir=None, backbone_only=False, opset_version=17, verbose=True, shape=None, batch_size=1, dynamic_batch=False, patch_size=None, format='onnx', quantization=None, calibration_data=None, max_images=100, *, backend=None, soc=None, fp16=True, max_batch_size=None, trt_metadata=False, trt_hardware_compatibility=None, trt_version_compatible=False, notes=None, output_name=None, trt_timing_cache=None, coreml_precision=None, coreml_neural_engine=False, coreai_precision=None, openvino_precision=None)

Export the trained model to ONNX, TFLite, TensorRT, ExecuTorch, CoreML, OpenVINO, or LiteRT format.

See the export documentation <https://rfdetr.roboflow.com/exports/>_ for more information.

Parameters:

Name Type Description Default

output_dir

str

Directory to write the exported model to.

'output'

infer_dir

str | None

Optional directory of sample images for dynamic-axes inference.

None

backbone_only

bool

Export the encoder and feature projectors without prediction heads. Returns one NCHW feature map per configured projector_scale level, followed by cross-attention feature levels when a separate cross-attention projector is present.

False

opset_version

int

ONNX opset version to target.

17

verbose

bool

Print export progress information.

True

shape

tuple[int, int] | None

(height, width) tuple; defaults to square at model resolution. Both dimensions must be divisible by patch_size * num_windows.

None

batch_size

int

Static batch size to bake into the ONNX graph; a positive integer. With dynamic_batch=True and format="tensorrt" it is also the batch the engine's optimization profile is tuned for.

1

dynamic_batch

bool

If True, export with a dynamic batch dimension so the model accepts variable batch sizes at runtime (spatial dimensions always stay fixed). Applies to the ONNX graph, and to format="tensorrt", where the engine is built with one optimization profile spanning batch 1 .. max_batch_size (tuned for batch_size); pass max_batch_size in that case. Not supported for ExecuTorch export on executorch 1.3.1 (raises NotImplementedError): the runtime cannot resize RF-DETR's windowed-attention reshapes, so a dynamic .pte runs only at the traced batch — export one .pte per batch size instead. Also unsupported for native CoreML (format="coreml"): fixed shapes are required for reliable ANE / GPU scheduling. Also refused for format="openvino", where it is not needed: omit dynamic_batch; the converted OpenVINO IR already accepts any batch size. Also refused for format="tflite": onnx2tf fails on the dynamic-batch graph (an Add against the dynamic-shaped encoder input reports mismatched dimensions) — export one .tflite per batch size instead. Also unsupported for format="litert": the .tflite bakes a fixed input shape; export one file per batch size instead.

False

patch_size

int | None

Backbone patch size. Defaults to the value stored in model_config.patch_size (typically 14 or 16). When provided explicitly it must match the instantiated model's patch size. Shape divisibility is validated against patch_size * num_windows.

None

format

str

Export format — "onnx" (default), "tflite", "tensorrt" (alias: "trt"), "executorch" (alias: "pte"), "coreml", "coreai", "openvino" or "litert", in any case ("ONNX" works too). "tflite" and "tensorrt" both first export to ONNX, then convert: "tflite" via onnx2tf (requires pip install rfdetr[tflite]); "tensorrt" via the TensorRT Python API (requires pip install rfdetr[tensorrt]). Unlike "onnx"/"tflite" portable serialization, "tensorrt" performs target-specific compilation at export time and by default produces a .trt engine tied to the build machine's GPU and TensorRT version; trt_hardware_compatibility and trt_version_compatible widen that. When "executorch" is selected the model is exported directly via torch.export to an ExecuTorch .pte file (no ONNX step), configured by backend / soc below. Requires pip install rfdetr[executorch]. "openvino" converts directly from PyTorch to OpenVINO IR format (requires pip install rfdetr[openvino]). "litert" converts directly from PyTorch to a LiteRT .tflite file via litert-torch (torch.export capture, no ONNX or TensorFlow step; requires pip install rfdetr[litert]) — the same file type as "tflite", reached without onnx2tf. It writes a single float32 graph with the two-stage query selection inside, so it runs on LiteRT's CPU (XNNPACK) delegate; keypoint models are refused before conversion, because litert-torch cannot lower the keypoint head's rank-4 matmul. When "coreml" is selected the model is exported via torch.export + coremltools to a native .mlpackage (no ONNX step; requires pip install rfdetr[coreml]). This is distinct from ExecuTorch's format="executorch", backend="coreml" path, which still produces a .pte. If you know that ExecuTorch delegate and expect format="coreml" to mean the same thing: it does not — pass format="executorch", backend="coreml" for the .pte route instead. Passing both format="coreml" and backend="coreml" together does not fall through to the ExecuTorch delegate; backend is ignored (with a warning) and the native .mlpackage path always runs. Detection, segmentation and keypoint models all have registry-clean and numerical-parity test coverage for this format (see tests/export/coreml/test_coreml_op_coverage.py / test_coreml_export.py). When "coreai" is selected the model is exported via torch.export + coreai-torch to an Apple Core AI .aimodel (requires pip install rfdetr[coreai]; the asset runs on iOS, iPadOS and macOS 27 or later).

.. warning:: TFLite, ExecuTorch, CoreML, and LiteRT export are experimental and subject to change; upstream dependency instabilities (onnx2tf, ai_edge_litert, executorch, coremltools, litert-torch) may affect results.

'onnx'

quantization

str | None

Quantization mode. Its meaning is per format, and formats not listed here ignore it (format="executorch", for example):

  • format="tflite" — one of None, "fp32", "fp16", "int8". None / "fp32" / "fp16" produce FP32 + FP16 .tflite files; "int8" additionally produces a dynamic-range INT8 model (INT8 weights, float activations; needs no calibration data).
  • format="onnx" — one of None, "fp32", "int8". "int8" writes a static QDQ model alongside the FP32 graph, with 8-bit weights and activations on the matrix multiplies only, and requires calibration_data. Attention scores, detection heads, normalization and the surrounding elementwise math stay in float; quantizing those costs accuracy and speed alike.
  • format="openvino" — one of None, "fp32", "int8". "int8" compresses the converted IR with NNCF (pip install nncf, not part of the rfdetr[openvino] extra) before it is written, so the .xml / .bin pair is already 8-bit, and requires calibration_data.
  • format="tensorrt" — None (an FP16 or FP32 engine, as fp16 says) or "int8": most of the backbone encoder and decoder run in INT8 and the rest in FP16 (placement in the TensorRT export guide), with activation ranges calibrated on calibration_data, which is then required (a detection model's backbone norms are smoothed and its backbone MLP biases corrected for INT8 on the same data). Detection, segmentation and keypoint models (the mask head and keypoint head stay FP16), static batch, fp16=True, TensorRT 10 or newer; the engine is named *_int8.trt.
  • format="litert" — accepts only None / "fp32" (it writes one float32 .tflite) and raises NotImplementedError for the other modes rather than silently ignoring them.
None

calibration_data

str | Path | ndarray[Any, Any] | None

Representative data for static INT8. Required for format="onnx", format="openvino" and format="tensorrt" with quantization="int8", where activation ranges are collected from it, and refused as missing there rather than defaulted. format="tensorrt" also refuses it without "int8". Accepts:

  • None — the default; no calibration data.
  • A directory path (str or :class:~pathlib.Path) containing images, preprocessed exactly as :meth:predict does.
  • A path (str or :class:~pathlib.Path) to a .npy file of shape (N, C, H, W), already normalized the way the model expects.
  • A :class:numpy.ndarray with the same format.

The resulting activation ranges decide the INT8 model's accuracy, so the data must be representative of the deployment domain — out-of-domain data yields a model that loads, runs, and is quietly wrong.

For format="tflite" it is ignored: any value other than None raises a UserWarning and is never read, so a missing path or a malformed array does not fail the export. No data could change the exported .tflite models: quantization="int8" produces a dynamic-range model whose weight scales come from the weights themselves, and fp32/fp16 involve no calibration.

None

max_images

int

Maximum number of images to load from a calibration_data directory (the first ones by file name) for ONNX / OpenVINO / TensorRT static INT8. Defaults to 100. Only used when calibration_data is a directory path; ignored for format="tflite" along with calibration_data.

100

backend

str | None

Hardware backend to specialize the export for. Required when format="executorch" and ignored — with a warning — for any other format. Accepted values for ExecuTorch: "xnnpack" (portable CPU, fp32), "coreml" (Apple devices, fp16; requires coremltools), and "qnn" (Qualcomm Snapdragon HTP, fp16; requires an ExecuTorch source build against the QAIRT SDK — not available via pip).

None

soc

str | None

Target SoC (System on Chip) — the specific Qualcomm Snapdragon chip the exported model will run on. Required when backend="qnn": the QNN backend compiles the .pte ahead-of-time for one chip's Hexagon Tensor Processor (HTP), unlike "xnnpack"/"coreml" which run on any device of their platform, so the target chip must be known at export time. Ignored — with a warning — for any other backend or format. Must be a :class:~executorch.backends.qualcomm.serialization.qc_schema.QcomChipset name, e.g. "SM8650" (Snapdragon 8 Gen 3); see that enum for the full list of supported chips. Has no effect for "xnnpack" or "coreml".

None

fp16

bool

Build the TensorRT engine with FP16 precision. Only applies when format="tensorrt" (alias "trt"); ignored for every other format. Defaults to True for lowest latency on NVIDIA GPUs. TensorRT 11+ removed the FP16 builder flag, so there the engine is built from an FP16-cast graph instead; engine I/O stays FP32 either way. A lean/partial TensorRT < 11 wheel that lacks the FP16 builder flag falls back to an FP32 engine with a warning instead — see :meth:~rfdetr.export._tensorrt.exporter.TensorRTExporter.build_engine for the full precision-resolution logic. Pass False for an FP32 engine.

True

max_batch_size

int | None

Largest batch a dynamic TensorRT engine must accept. Only read when format="tensorrt" and dynamic_batch=True, where it is required: the engine is built with one optimization profile spanning batch 1 .. max_batch_size and tuned for batch_size (batch_size <= max_batch_size). Ignored for every other format or combination; passing a non-None value there emits a UserWarning instead of silently doing nothing.

None

trt_metadata

bool

Also write <engine>.json beside a format="tensorrt" engine, so a consumer that does not import the model (a C++ service, Triton, DeepStream) can learn the engine's input size and normalization, output names, batch profile, the precision actually built, the TensorRT version and GPU it was built on, and the engine file's size and SHA-256. A .trt file has no slot for this. It does not record class names. False (default) writes no description. Ignored for every other format; True there emits a UserWarning instead of silently doing nothing.

False

trt_hardware_compatibility

Literal['ampere_plus', 'same_compute_capability'] | None

Ask TensorRT for a format="tensorrt" engine that other GPUs may run too. "ampere_plus" targets NVIDIA Ampere GPUs (compute capability 8.x) and newer, and needs an Ampere or newer GPU to build; "same_compute_capability" targets GPUs that share the building GPU's compute capability. Not supported on Jetson (JetPack) or DriveOS. The engine can run slower than one built for a single GPU. None (default) builds for the building GPU only. Ignored for every other format; passing a non-None value there emits a UserWarning instead of silently doing nothing.

None

trt_version_compatible

bool

Ask TensorRT for a format="tensorrt" engine that other releases of the same TensorRT major version may load. It worked between TensorRT 11.2 and 11.3, in both directions; an engine did not load across major versions (10 and 11), nor between 10.13 and 10.16. It needs TensorRT's lean runtime library, a separate package from tensorrt (tensorrt-lean-cu*-libs). The engine built by TensorRT 11 carries host code: load it with TRTInference(..., engine_host_code_allowed=True), and only from a file you trust. False (default) builds an engine that loads on the building TensorRT version only. Ignored for every other format; True there emits a UserWarning instead of silently doing nothing.

False

trt_timing_cache

str | PathLike[str] | None

File in which TensorRT keeps the kernel timings it measures while building an engine, for format="tensorrt". A build loads the file when it exists and writes the merged timings back, so a later build with the same layer shapes at the same precision and batch profile, on the same GPU and TensorRT version, skips the search it already did. The timings depend on the layers' shapes, not the weights, so a re-export with new weights reuses them for every layer whose shape did not change. A cache from a matching dynamic-batch build helps too; reuse between a static and a dynamic profile, or across precisions, was not measured. A relative path is relative to the working directory, not to output_dir. A cache written by another TensorRT major version, or an empty or damaged file, does not stop the build: TensorRT logs an error, builds as if there were no cache, and the file gets this build's timings. A failed cache write does fail the export, and the built engine is not saved (Polygraphy writes the cache before the engine is returned); the engine can simply be rebuilt. None (default) reads and writes no file. Ignored for every other format; passing a non-None value there emits a UserWarning instead of silently doing nothing.

None

notes

object

Optional user-defined metadata (string, dict, list, or any JSON-serialisable value) to embed in the metadata of the exported artifact, for a format that carries it (see "Read Embedded Notes" in the export documentation for where each format stores it). When None no metadata entry is written. String values are stored verbatim; all other types are JSON-encoded so consumers must call json.loads() to recover a dict or list. The same value can be passed to :meth:train so the checkpoint and the exported artifact share the same provenance information. Ignored for format="executorch", format="openvino", and format="litert": those artifacts have no ONNX-style metadata slot, and a non-None value emits a UserWarning instead of being embedded. Every other format refuses a value JSON cannot encode (NaN, Infinity, an arbitrary object) before any model work.

None

coreml_precision

str | None

ct.convert compute precision for format="coreml" — None (default) or "float32" selects FP32 (tight CPU parity with eager PyTorch); "float16" selects a smaller ANE-oriented bundle (expect larger numeric drift) whose outputs are still float32. Bundles declare iOS 15 / macOS 12 as the minimum OS. Ignored for every other format.

None

coreml_neural_engine

bool

For format="coreml", export a graph for the Apple Neural Engine, with the same weights: the backbone attention runs as one einsum pair per head in query chunks, and the two-stage query selection as a one-hot matmul. In an fp16 export, only the deformable-attention resample and its cast ops run on the CPU. This graph is faster than the default graph on the Neural Engine and slower on the CPU and the GPU, so use it only with CPU_AND_NE or ALL compute units. Ignored for every other format.

False

coreai_precision

str | None

Precision the graph is traced and stored in for format="coreai" — None (default) or "float32", or "float16" for a half-size asset whose input and outputs are float16 too. A float16 keypoint model warns: the asset aborts the process on the Neural Engine. Ignored for every other format.

None

openvino_precision

str | None

"float32", "float16", or None (default) for format="openvino" — None keeps OpenVINO's own compress_to_fp16=True default; "float32" disables FP16 weight compression, controlling IR storage precision only (execution precision still depends on the compiled device — not guaranteed to match eager PyTorch on non-CPU devices). Ignored for every other format.

None

output_name

str | None

Full filename override (without extension), e.g. "my-model". When set, takes precedence over the model's variant name (self.size) and the exported file is named {output_name}.{ext} verbatim — this also suppresses the _fp32/_fp16/_{backend} detail suffix that would otherwise be appended to encode the resolved precision/backend/SoC (see format / coreml_precision / backend / soc / fp16 above), and the TensorRT portability suffixes (_ampere_plus, _same_compute_capability, _version_compatible). Sanitized against path traversal: only the basename is used, with an export file extension such as .onnx stripped; other dots, as in "my-model-v1.2", are kept. Exception: format="tflite" always writes multiple files (one per precision/quantization mode), so the _fp32/_fp16/ _dynamic_range_quant suffix is unavoidable even with output_name set — it becomes the stem instead of the model's variant name. Exceptions: ONNX, CoreML, ExecuTorch, OpenVINO, LiteRT, and TensorRT with backbone_only=True append -backbone before the extension (e.g., {output_name}-backbone.onnx); TFLite writes per-precision files instead of a single {output_name}.tflite file. TFLite filenames may include a _gs_patched infix before the precision suffix when GridSample ops are patched, e.g. {output_name}_gs_patched_fp32.tflite; this is the standard RF-DETR path.

None

Returns:

Type Description
Path

Path to the exported model file (.onnx, .tflite for both TFLite and LiteRT, .trt,

Path

.pte, .mlpackage, .aimodel or .xml for OpenVINO). With format="onnx" and

Path

quantization="int8" it is the quantized {stem}_int8.onnx; the FP32 {stem}.onnx it was derived

Path

from is left beside it.

Raises:

Type Description
ValueError

If format is unrecognized; if batch_size, or max_batch_size when given, is not a positive integer; if format="executorch" and backend is missing, unrecognized, or (for backend="qnn") soc is missing; if the resolved export shape is not divisible by patch_size * num_windows; if coreml_precision/coreai_precision/openvino_precision, or quantization for format="tflite", is not one of their accepted values; if notes holds a non-finite float or a circular reference, for a format that embeds it; or if format="tensorrt" with dynamic_batch=True lacks max_batch_size or has batch_size > max_batch_size; if quantization is not a recognized mode for format="onnx" or format="openvino" (only None, "fp32" and "int8" are accepted there); if quantization="int8" is combined with format="onnx" or format="openvino" and no calibration_data; or, for those two INT8 requests, if calibration_data points to a missing path, a file that is not .npy, an image-less directory, or an array that is not (N, C, H, W) or does not match the graph's input size. Also raised for format="tensorrt" when quantization is neither None nor "int8", calibration_data is given without "int8", "int8" has no calibration_data, max_images is not a positive integer, or "int8" is combined with fp16=False, dynamic_batch=True or backbone_only=True; calibration_data that is not a directory, a .npy path or a non-empty (N, C, H, W) float array, a path that does not exist, a directory without an image, a file that is not .npy and a .npy that does not hold such an array are all refused before the forward pass. During the INT8 conversion, calibration data that yields no usable image or a non-finite range, a range too large for an FP16 scale, and attention INT8 cannot be placed around, also raise. Also raised for format="tensorrt" when trt_metadata is not a bool. Also raised for format="tensorrt" when trt_hardware_compatibility is neither None, "ampere_plus" nor "same_compute_capability", when trt_version_compatible is not a bool, when the installed TensorRT has no hardware compatibility level of the requested name, or when trt_hardware_compatibility="ampere_plus" is asked of a CUDA device older than Ampere. Also raised for format="tensorrt" when trt_timing_cache is not a non-empty file path, ends in a path separator, or is a directory.

TypeError

If notes holds a value JSON cannot encode, for a format that embeds it.

OSError

If format="tensorrt" and the directory or files of trt_timing_cache cannot be created or written, or it is a symbolic link to a missing file.

NotImplementedError

If dynamic_batch=True is combined with format="executorch", format="coreml", format="openvino", format="tflite", or format="litert" — those paths require a fixed batch size; if format="litert" is combined with a quantization other than None / "fp32"; if format="litert" is asked to export a keypoint model (backbone_only=True still converts).

ImportError

If the optional dependencies for the requested format/backend are not installed (e.g. rfdetr[onnx], rfdetr[tensorrt], rfdetr[executorch], rfdetr[coreml], coremltools for ExecuTorch backend="coreml", openvino for OpenVINO export, rfdetr[litert] for LiteRT export, onnxruntime for format="onnx" with quantization="int8", nncf for format="openvino" with quantization="int8", or an ExecuTorch source build against the QAIRT SDK for backend="qnn"); also raised for format="tensorrt" with fp16=True on a strongly typed TensorRT (11+) if onnx/onnxconverter-common are not installed to cast the graph — install rfdetr[tensorrt] for the complete set, or pass fp16=False; and for format="tensorrt" with quantization="int8" without onnxruntime or onnxconverter-common, or on a TensorRT older than 10. Each format's availability check runs before the model does, and so do the check for TensorRT's lean runtime library that trt_version_compatible=True needs and the TensorRT INT8 host checks; what they do not cover (a backend's extension, the TensorRT cast's packages, the Core AI runtime package, onnxruntime or nncf for INT8) is found missing only during the conversion.

RuntimeError

If called after the model has undergone in-place inference optimization (the original model has been cleared; instantiate a new :class:RFDETR to export).

OSError

If trt_metadata=True and the description file cannot be written, after the engine was built.

Examples:

Skipped because this example requires pretrained weights and optional ONNX export dependencies.

>>> from rfdetr import RFDETRSmall
>>> model = RFDETRSmall()
>>> model.export(format="onnx", output_dir="output")
Source code in src/rfdetr/detr.py
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
def export(
    self,
    output_dir: str = "output",
    infer_dir: str | None = None,
    backbone_only: bool = False,
    opset_version: int = 17,
    verbose: bool = True,
    shape: tuple[int, int] | None = None,
    batch_size: int = 1,
    dynamic_batch: bool = False,
    patch_size: int | None = None,
    format: str = "onnx",
    quantization: str | None = None,
    calibration_data: str | Path | np.ndarray[Any, Any] | None = None,
    max_images: int = 100,
    *,
    backend: str | None = None,
    soc: str | None = None,
    fp16: bool = True,
    max_batch_size: int | None = None,
    trt_metadata: bool = False,
    trt_hardware_compatibility: Literal["ampere_plus", "same_compute_capability"] | None = None,
    trt_version_compatible: bool = False,
    notes: object = None,
    output_name: str | None = None,
    trt_timing_cache: str | os.PathLike[str] | None = None,
    coreml_precision: str | None = None,
    coreml_neural_engine: bool = False,
    coreai_precision: str | None = None,
    openvino_precision: str | None = None,
) -> Path:
    """Export the trained model to ONNX, TFLite, TensorRT, ExecuTorch, CoreML, OpenVINO, or LiteRT format.

    See the `export documentation <https://rfdetr.roboflow.com/exports/>`_ for more information.

    Args:
        output_dir: Directory to write the exported model to.
        infer_dir: Optional directory of sample images for dynamic-axes inference.
        backbone_only: Export the encoder and feature projectors without prediction heads. Returns one NCHW
            feature map per configured ``projector_scale`` level, followed by cross-attention feature levels
            when a separate cross-attention projector is present.
        opset_version: ONNX opset version to target.
        verbose: Print export progress information.
        shape: ``(height, width)`` tuple; defaults to square at model resolution.
            Both dimensions must be divisible by ``patch_size * num_windows``.
        batch_size: Static batch size to bake into the ONNX graph; a positive integer. With ``dynamic_batch=True``
            and ``format="tensorrt"`` it is also the batch the engine's optimization profile is tuned for.
        dynamic_batch: If True, export with a dynamic batch dimension
            so the model accepts variable batch sizes at runtime
            (spatial dimensions always stay fixed). Applies to the ONNX
            graph, and to ``format="tensorrt"``, where the engine is built with one
            optimization profile spanning batch ``1 .. max_batch_size`` (tuned for *batch_size*); pass
            *max_batch_size* in that case. Not supported for ExecuTorch export on
            executorch 1.3.1 (raises ``NotImplementedError``): the runtime
            cannot resize RF-DETR's windowed-attention reshapes, so a
            dynamic ``.pte`` runs only at the traced batch — export one
            ``.pte`` per batch size instead. Also unsupported for native CoreML
            (``format="coreml"``): fixed shapes are required for reliable ANE / GPU scheduling.
            Also refused for ``format="openvino"``, where it is not needed: omit ``dynamic_batch``;
            the converted OpenVINO IR already accepts any batch size.
            Also refused for ``format="tflite"``: ``onnx2tf`` fails on the dynamic-batch graph (an ``Add``
            against the dynamic-shaped encoder input reports mismatched dimensions) — export one ``.tflite``
            per batch size instead. Also unsupported for ``format="litert"``: the ``.tflite`` bakes a fixed
            input shape; export one file per batch size instead.
        patch_size: Backbone patch size. Defaults to the value stored
            in ``model_config.patch_size`` (typically 14 or 16). When
            provided explicitly it must match the instantiated model's
            patch size. Shape divisibility is validated against
            ``patch_size * num_windows``.
        format: Export format — ``"onnx"`` (default), ``"tflite"``, ``"tensorrt"`` (alias: ``"trt"``),
            ``"executorch"`` (alias: ``"pte"``), ``"coreml"``, ``"coreai"``, ``"openvino"`` or ``"litert"``,
            in any case (``"ONNX"`` works too).
            ``"tflite"`` and ``"tensorrt"`` both first export to ONNX,
            then convert: ``"tflite"`` via ``onnx2tf`` (requires
            ``pip install rfdetr[tflite]``); ``"tensorrt"`` via the
            TensorRT Python API (requires ``pip install rfdetr[tensorrt]``).
            Unlike ``"onnx"``/``"tflite"`` portable serialization,
            ``"tensorrt"`` performs target-specific compilation at
            export time and by default produces a ``.trt`` engine tied
            to the build machine's GPU and TensorRT version;
            ``trt_hardware_compatibility`` and ``trt_version_compatible``
            widen that.
            When ``"executorch"`` is selected the model is exported
            directly via ``torch.export`` to an ExecuTorch
            ``.pte`` file (no ONNX step), configured by *backend* / *soc* below.  Requires
            ``pip install rfdetr[executorch]``. ``"openvino"`` converts directly from PyTorch to OpenVINO IR
            format (requires ``pip install rfdetr[openvino]``). ``"litert"`` converts directly from PyTorch to
            a LiteRT ``.tflite`` file via ``litert-torch`` (``torch.export`` capture, no ONNX or TensorFlow
            step; requires ``pip install rfdetr[litert]``) — the same file type as ``"tflite"``, reached
            without ``onnx2tf``. It writes a single float32 graph with the two-stage query selection
            inside, so it runs on LiteRT's CPU (XNNPACK) delegate; keypoint models are refused before
            conversion, because litert-torch cannot lower the keypoint head's rank-4 matmul.
            When ``"coreml"`` is selected the model is exported via ``torch.export`` + ``coremltools`` to a
            native ``.mlpackage`` (no ONNX step; requires
            ``pip install rfdetr[coreml]``). This is distinct from
            ExecuTorch's ``format="executorch", backend="coreml"``
            path, which still produces a ``.pte``. If you know that
            ExecuTorch delegate and expect ``format="coreml"`` to mean
            the same thing: it does not — pass
            ``format="executorch", backend="coreml"`` for the ``.pte``
            route instead. Passing both ``format="coreml"`` and
            ``backend="coreml"`` together does **not** fall through
            to the ExecuTorch delegate; ``backend`` is ignored (with
            a warning) and the native ``.mlpackage`` path always runs.
            Detection, segmentation and keypoint models all have registry-clean and numerical-parity
            test coverage for this format (see ``tests/export/coreml/test_coreml_op_coverage.py`` /
            ``test_coreml_export.py``).
            When ``"coreai"`` is selected the model is exported via ``torch.export`` + ``coreai-torch`` to an
            Apple Core AI ``.aimodel`` (requires ``pip install rfdetr[coreai]``; the asset runs on iOS,
            iPadOS and macOS 27 or later).

            .. warning::
                TFLite, ExecuTorch, CoreML, and LiteRT export are experimental
                and subject to change; upstream dependency instabilities
                (``onnx2tf``, ``ai_edge_litert``, ``executorch``,
                ``coremltools``, ``litert-torch``) may affect results.
        quantization: Quantization mode.  Its meaning is per format, and formats not listed here ignore it
            (``format="executorch"``, for example):

            * ``format="tflite"`` — one of ``None``, ``"fp32"``, ``"fp16"``, ``"int8"``.  ``None`` / ``"fp32"``
              / ``"fp16"`` produce FP32 + FP16 ``.tflite`` files; ``"int8"`` additionally produces a
              **dynamic-range** INT8 model (INT8 weights, float activations; needs no calibration data).
            * ``format="onnx"`` — one of ``None``, ``"fp32"``, ``"int8"``.  ``"int8"`` writes a **static** QDQ
              model alongside the FP32 graph, with 8-bit weights *and* activations on the matrix multiplies
              only, and **requires** *calibration_data*.  Attention scores, detection heads, normalization and
              the surrounding elementwise math stay in float; quantizing those costs accuracy and speed alike.
            * ``format="openvino"`` — one of ``None``, ``"fp32"``, ``"int8"``.  ``"int8"`` compresses the
              converted IR with NNCF (``pip install nncf``, not part of the ``rfdetr[openvino]`` extra) before it
              is written, so the ``.xml`` / ``.bin`` pair is already 8-bit, and **requires** *calibration_data*.
            * ``format="tensorrt"`` — ``None`` (an FP16 or FP32 engine, as ``fp16`` says) or ``"int8"``: most of
              the backbone encoder and decoder run in INT8 and the rest in FP16 (placement in the TensorRT export
              guide), with activation ranges calibrated on *calibration_data*, which is then required (a
              detection model's backbone norms are smoothed and its backbone MLP biases corrected for INT8 on the
              same data).  Detection, segmentation and keypoint models (the mask head and keypoint head stay
              FP16), static batch, ``fp16=True``, TensorRT 10 or newer; the engine is named ``*_int8.trt``.
            * ``format="litert"`` — accepts only ``None`` / ``"fp32"`` (it writes one float32 ``.tflite``) and
              raises ``NotImplementedError`` for the other modes rather than silently ignoring them.
        calibration_data: Representative data for static INT8.  **Required** for ``format="onnx"``,
            ``format="openvino"`` and ``format="tensorrt"`` with ``quantization="int8"``, where activation ranges
            are collected from it, and refused as missing there rather than defaulted.  ``format="tensorrt"``
            also refuses it without ``"int8"``.  Accepts:

            * ``None`` — the default; no calibration data.
            * A **directory path** (``str`` or :class:`~pathlib.Path`) containing images, preprocessed exactly as
              :meth:`predict` does.
            * A path (``str`` or :class:`~pathlib.Path`) to a ``.npy`` file of shape ``(N, C, H, W)``, already
              normalized the way the model expects.
            * A :class:`numpy.ndarray` with the same format.

            The resulting activation ranges decide the INT8 model's accuracy, so the data must be representative of
            the deployment domain — out-of-domain data yields a model that loads, runs, and is quietly wrong.

            For ``format="tflite"`` it is ignored: any value other than ``None`` raises a ``UserWarning`` and is
            never read, so a missing path or a malformed array does not fail the export. No data could change the
            exported ``.tflite`` models: ``quantization="int8"`` produces a dynamic-range model whose weight
            scales come from the weights themselves, and fp32/fp16 involve no calibration.
        max_images: Maximum number of images to load from a *calibration_data* directory (the first ones by file
            name) for ONNX / OpenVINO / TensorRT static INT8.  Defaults to ``100``.  Only used when
            *calibration_data* is a directory path; ignored for ``format="tflite"`` along with *calibration_data*.
        backend: Hardware backend to specialize the export for.  Required when ``format="executorch"`` and
            ignored — with a warning — for any other format.  Accepted values for ExecuTorch:
            ``"xnnpack"`` (portable CPU, fp32), ``"coreml"`` (Apple devices, fp16; requires ``coremltools``),
            and ``"qnn"`` (Qualcomm Snapdragon HTP, fp16; requires an ExecuTorch source build against the
            QAIRT SDK — not available via pip).
        soc: Target SoC (System on Chip) — the specific Qualcomm Snapdragon chip the exported model will run
            on.  Required when ``backend="qnn"``: the QNN backend compiles the ``.pte`` ahead-of-time for one
            chip's Hexagon Tensor Processor (HTP), unlike ``"xnnpack"``/``"coreml"`` which run on any device of
            their platform, so the target chip must be known at export time.  Ignored — with a warning — for
            any other backend or format.  Must be a
            :class:`~executorch.backends.qualcomm.serialization.qc_schema.QcomChipset` name, e.g. ``"SM8650"``
            (Snapdragon 8 Gen 3); see that enum for the full list of supported chips.  Has no effect for
            ``"xnnpack"`` or ``"coreml"``.
        fp16: Build the TensorRT engine with FP16 precision.  Only applies when ``format="tensorrt"``
            (alias ``"trt"``); ignored for every other format.  Defaults to ``True`` for lowest latency
            on NVIDIA GPUs.  TensorRT 11+ removed the FP16 builder flag, so there the engine is built
            from an FP16-cast graph instead; engine I/O stays FP32 either way.  A lean/partial
            TensorRT < 11 wheel that lacks the FP16 builder flag falls back to an FP32 engine with a
            warning instead — see :meth:`~rfdetr.export._tensorrt.exporter.TensorRTExporter.build_engine`
            for the full precision-resolution logic.  Pass ``False`` for an FP32 engine.
        max_batch_size: Largest batch a dynamic TensorRT engine must accept.  Only read when
            ``format="tensorrt"`` and ``dynamic_batch=True``, where it is required: the engine is built with
            one optimization profile spanning batch ``1 .. max_batch_size`` and tuned for *batch_size*
            (``batch_size <= max_batch_size``).  Ignored for every other format or combination; passing a
            non-``None`` value there emits a ``UserWarning`` instead of silently doing nothing.
        trt_metadata: Also write ``<engine>.json`` beside a ``format="tensorrt"`` engine, so a consumer that does
            not import the model (a C++ service, Triton, DeepStream) can learn the engine's input size and
            normalization, output names, batch profile, the precision actually built, the TensorRT version and GPU
            it was built on, and the engine file's size and SHA-256.  A ``.trt`` file has no slot for this.  It does
            not record class names.
            ``False`` (default) writes no description.  Ignored for every other format; ``True`` there emits a
            ``UserWarning`` instead of silently doing nothing.
        trt_hardware_compatibility: Ask TensorRT for a ``format="tensorrt"`` engine that other GPUs may run too.
            ``"ampere_plus"`` targets NVIDIA Ampere GPUs (compute capability 8.x) and newer, and needs an Ampere or
            newer GPU to build; ``"same_compute_capability"`` targets GPUs that share the building GPU's compute
            capability.  Not supported on Jetson (JetPack) or DriveOS.  The engine can run slower than one built
            for a single GPU.  ``None`` (default) builds for the building GPU only.  Ignored for every other
            format; passing a non-``None`` value there emits a ``UserWarning`` instead of silently doing nothing.
        trt_version_compatible: Ask TensorRT for a ``format="tensorrt"`` engine that other releases of the same
            TensorRT major version may load.  It worked between TensorRT 11.2 and 11.3, in both directions; an
            engine did not load across major versions (10 and 11), nor between 10.13 and 10.16.  It needs
            TensorRT's lean runtime library, a separate package from ``tensorrt`` (``tensorrt-lean-cu*-libs``).
            The engine built by TensorRT 11 carries host code: load it with
            ``TRTInference(..., engine_host_code_allowed=True)``, and only from a file you trust.  ``False``
            (default) builds an engine that loads on the building TensorRT version only.  Ignored for every other
            format; ``True`` there emits a ``UserWarning`` instead of silently doing nothing.
        trt_timing_cache: File in which TensorRT keeps the kernel timings it measures while building an engine, for
            ``format="tensorrt"``.  A build loads the file when it exists and writes the merged timings back, so a
            later build with the same layer shapes at the same precision and batch profile, on the same GPU and
            TensorRT version, skips the search it already did.  The timings depend on the layers' shapes, not the
            weights, so a re-export with new weights reuses them for every layer whose shape did not change.  A
            cache from a matching dynamic-batch build helps too; reuse between a static and a dynamic profile, or
            across precisions, was not measured.  A relative path is relative to the working directory, not to
            *output_dir*.  A cache written by another TensorRT major version, or an empty or damaged file, does
            not stop the build: TensorRT logs an error, builds as if there were no cache, and the file gets this
            build's timings.  A failed cache write does fail the export, and the built engine is not saved
            (Polygraphy writes the cache before the engine is returned); the engine can simply be rebuilt.
            ``None`` (default) reads and writes no file.  Ignored for every other format; passing a non-``None``
            value there emits a ``UserWarning`` instead of silently doing nothing.
        notes: Optional user-defined metadata (string, dict, list,
            or any JSON-serialisable value) to embed in the metadata of
            the exported artifact, for a format that carries it (see
            "Read Embedded Notes" in the export documentation for where
            each format stores it).
            When ``None`` no metadata entry is written. String values
            are stored verbatim; all other types are JSON-encoded so
            consumers must call ``json.loads()`` to recover a dict or
            list. The same value can be passed to :meth:`train` so the
            checkpoint and the exported artifact share the same provenance
            information.
            **Ignored for ``format="executorch"``, ``format="openvino"``, and ``format="litert"``**: those
            artifacts have no ONNX-style metadata slot, and a non-``None`` value emits a ``UserWarning``
            instead of being embedded. Every other format refuses a value JSON cannot encode (``NaN``,
            ``Infinity``, an arbitrary object) before any model work.
        coreml_precision: ``ct.convert`` compute precision for ``format="coreml"`` — ``None`` (default) or
            ``"float32"`` selects FP32 (tight CPU parity with eager
            PyTorch); ``"float16"`` selects a smaller
            ANE-oriented bundle (expect larger numeric drift) whose outputs are still float32. Bundles
            declare iOS 15 / macOS 12 as the minimum OS. Ignored for every other format.
        coreml_neural_engine: For ``format="coreml"``, export a graph for the Apple Neural Engine, with the same
            weights: the backbone attention runs as one einsum pair per head in query chunks, and the two-stage
            query selection as a one-hot matmul. In an fp16 export, only the deformable-attention ``resample`` and
            its ``cast`` ops run on the CPU. This graph is faster than the default graph on the Neural Engine and
            slower on the CPU and the GPU, so use it only with ``CPU_AND_NE`` or ``ALL`` compute units. Ignored for
            every other format.
        coreai_precision: Precision the graph is traced and stored in for ``format="coreai"`` — ``None``
            (default) or ``"float32"``, or ``"float16"`` for a half-size asset whose input and outputs are
            float16 too. A float16 keypoint model warns: the asset aborts the process on the Neural Engine.
            Ignored for every other format.
        openvino_precision: ``"float32"``, ``"float16"``, or ``None`` (default) for ``format="openvino"``
            — ``None`` keeps OpenVINO's own ``compress_to_fp16=True`` default; ``"float32"`` disables
            FP16 weight compression, controlling IR *storage* precision only (execution precision still
            depends on the compiled device — not guaranteed to match eager PyTorch on non-CPU devices).
            Ignored for every other format.
        output_name: Full filename override (without extension), e.g. ``"my-model"``. When set, takes
            precedence over the model's variant name (``self.size``) and the exported file is named
            ``{output_name}.{ext}`` verbatim — this also suppresses the ``_fp32``/``_fp16``/``_{backend}``
            detail suffix that would otherwise be appended to encode the resolved precision/backend/SoC
            (see *format* / *coreml_precision* / *backend* / *soc* / *fp16* above), and the TensorRT
            portability suffixes (``_ampere_plus``, ``_same_compute_capability``, ``_version_compatible``).
            Sanitized against path traversal: only the basename is used, with an export file extension such
            as ``.onnx`` stripped; other dots, as in ``"my-model-v1.2"``, are kept. Exception: ``format="tflite"``
            always writes multiple files (one per precision/quantization mode), so the ``_fp32``/``_fp16``/
            ``_dynamic_range_quant`` suffix is unavoidable even with
            *output_name* set — it becomes the stem
            instead of the model's variant name.
            Exceptions: ONNX, CoreML, ExecuTorch, OpenVINO, LiteRT, and TensorRT with ``backbone_only=True``
            append ``-backbone`` before the extension (e.g., ``{output_name}-backbone.onnx``); TFLite writes
            per-precision files
            instead of a single ``{output_name}.tflite`` file. TFLite filenames may include a ``_gs_patched`` infix
            before the precision suffix when GridSample ops are patched, e.g.
            ``{output_name}_gs_patched_fp32.tflite``; this is the standard RF-DETR path.

    Returns:
        Path to the exported model file (``.onnx``, ``.tflite`` for both TFLite and LiteRT, ``.trt``,
        ``.pte``, ``.mlpackage``, ``.aimodel`` or ``.xml`` for OpenVINO).  With ``format="onnx"`` and
        ``quantization="int8"`` it is the quantized ``{stem}_int8.onnx``; the FP32 ``{stem}.onnx`` it was derived
        from is left beside it.

    Raises:
        ValueError: If ``format`` is unrecognized; if ``batch_size``, or ``max_batch_size`` when given, is not
            a positive integer; if
            ``format="executorch"`` and ``backend`` is missing, unrecognized, or (for ``backend="qnn"``)
            ``soc`` is missing; if the resolved export shape is not divisible by ``patch_size * num_windows``;
            if ``coreml_precision``/``coreai_precision``/``openvino_precision``, or ``quantization`` for
            ``format="tflite"``, is not one of their accepted values; if ``notes`` holds a non-finite float or a
            circular reference, for a format that embeds it; or if ``format="tensorrt"`` with
            ``dynamic_batch=True`` lacks ``max_batch_size`` or has ``batch_size > max_batch_size``; if
            ``quantization`` is not a recognized mode for ``format="onnx"`` or ``format="openvino"`` (only
            ``None``, ``"fp32"`` and ``"int8"`` are accepted there); if ``quantization="int8"`` is combined with
            ``format="onnx"`` or ``format="openvino"`` and no ``calibration_data``; or, for those two INT8
            requests, if ``calibration_data`` points to a missing path, a file that is not ``.npy``, an
            image-less directory, or an array that is not ``(N, C, H, W)`` or does not match the graph's
            input size.
            Also raised for ``format="tensorrt"`` when ``quantization`` is neither ``None`` nor ``"int8"``,
            ``calibration_data`` is given without ``"int8"``, ``"int8"`` has no ``calibration_data``,
            ``max_images`` is not a positive integer, or ``"int8"`` is combined with ``fp16=False``,
            ``dynamic_batch=True`` or ``backbone_only=True``; ``calibration_data`` that is not a directory, a
            ``.npy`` path or a non-empty ``(N, C, H, W)`` float array, a path that does not exist, a directory
            without an image, a file that is not ``.npy`` and a ``.npy`` that does not hold such an array are all
            refused before the forward pass. During the INT8 conversion, calibration data that yields no usable
            image or a non-finite range, a range too large for an FP16 scale, and attention INT8 cannot be placed
            around, also raise.
            Also raised for ``format="tensorrt"`` when ``trt_metadata`` is not a ``bool``.
            Also raised for ``format="tensorrt"`` when ``trt_hardware_compatibility`` is neither ``None``,
            ``"ampere_plus"`` nor ``"same_compute_capability"``, when ``trt_version_compatible`` is not a
            ``bool``, when the installed TensorRT has no hardware compatibility level of the requested name, or
            when ``trt_hardware_compatibility="ampere_plus"`` is asked of a CUDA device older than Ampere.
            Also raised for ``format="tensorrt"`` when ``trt_timing_cache`` is not a non-empty file path, ends in
            a path separator, or is a directory.
        TypeError: If ``notes`` holds a value JSON cannot encode, for a format that embeds it.
        OSError: If ``format="tensorrt"`` and the directory or files of ``trt_timing_cache`` cannot be created or
            written, or it is a symbolic link to a missing file.
        NotImplementedError: If ``dynamic_batch=True`` is combined with ``format="executorch"``,
            ``format="coreml"``, ``format="openvino"``, ``format="tflite"``, or ``format="litert"`` — those
            paths require a fixed batch size; if ``format="litert"`` is combined with a ``quantization``
            other than ``None`` / ``"fp32"``; if ``format="litert"`` is asked to export a keypoint model
            (``backbone_only=True`` still converts).
        ImportError: If the optional dependencies for the requested
            ``format``/``backend`` are not installed (e.g.
            ``rfdetr[onnx]``, ``rfdetr[tensorrt]``, ``rfdetr[executorch]``,
            ``rfdetr[coreml]``, ``coremltools`` for ExecuTorch
            ``backend="coreml"``, ``openvino`` for OpenVINO export,
            ``rfdetr[litert]`` for LiteRT export,
            ``onnxruntime`` for ``format="onnx"`` with ``quantization="int8"``, ``nncf`` for
            ``format="openvino"`` with ``quantization="int8"``,
            or an ExecuTorch source build against the QAIRT SDK for
            ``backend="qnn"``); also raised for ``format="tensorrt"`` with ``fp16=True`` on a
            strongly typed TensorRT (11+) if ``onnx``/``onnxconverter-common`` are not installed
            to cast the graph — install ``rfdetr[tensorrt]`` for the complete set, or pass
            ``fp16=False``; and for ``format="tensorrt"`` with ``quantization="int8"`` without onnxruntime or
            ``onnxconverter-common``, or on a TensorRT older than 10. Each format's availability check runs before
            the model does, and so do the check for TensorRT's lean runtime library that
            ``trt_version_compatible=True`` needs and the TensorRT INT8 host checks; what they do not cover (a
            backend's extension, the TensorRT cast's packages, the Core AI runtime package, ``onnxruntime`` or
            ``nncf`` for INT8) is found missing only during the conversion.
        RuntimeError: If called after the model has undergone in-place inference optimization (the original
            model has been cleared; instantiate a new :class:`RFDETR` to export).
        OSError: If ``trt_metadata=True`` and the description file cannot be written, after the engine was built.

    Examples:
        Skipped because this example requires pretrained weights and optional ONNX export dependencies.

        >>> from rfdetr import RFDETRSmall  # doctest: +SKIP
        >>> model = RFDETRSmall()  # doctest: +SKIP
        >>> model.export(format="onnx", output_dir="output")  # doctest: +SKIP
    """
    from rfdetr.export._backend import _resolve_export_backend
    from rfdetr.export.base import reject_unsupported_dynamic_batch
    from rfdetr.export.prepare import prepare_export_graph, validate_export_shape
    from rfdetr.export.registry import resolve_exporter

    # Every format builds its example batch from this, so it is checked before any exporter is imported.
    format, export_batch_size, export_max_batch_size = RFDETR._resolve_export_batch(
        format, batch_size, dynamic_batch, max_batch_size
    )
    for keyword, requested in (
        ("trt_metadata", trt_metadata),
        ("trt_hardware_compatibility", trt_hardware_compatibility is not None),
        ("trt_version_compatible", trt_version_compatible),
        ("trt_timing_cache", trt_timing_cache is not None),
    ):
        if requested and format != "tensorrt":
            warnings.warn(
                f"`{keyword}` is only used for format='tensorrt' (got format={format!r}). "
                "This argument is ignored.",
                UserWarning,
                stacklevel=2,
            )
    backend, soc = _resolve_export_backend(format, backend, soc)
    # Refuse a statically impossible request from the registry's own capability data, before resolving the
    # exporter imports the format's heavy optional dependency (coremltools, executorch, openvino, ...) and long
    # before the user pays for a full DINOv2 forward pass.
    reject_unsupported_dynamic_batch(format, dynamic_batch=dynamic_batch)
    if getattr(self, "_optimized_inplace", False) or self.model.model is None:
        raise RuntimeError(
            "RFDETR.export() is not available after inplace optimization. "
            "The original model has been cleared. Create a new RFDETR instance."
        )
    # The shape needs only the model's configuration, so it is checked with the arguments above, before the exporter
    # is resolved.
    shape = validate_export_shape(
        shape,
        patch_size,
        self.model_config,
        self.model.resolution,
        resolve_patch_size=_resolve_patch_size,
        validate_shape_dims=_validate_shape_dims,
    )
    exporter_class = resolve_exporter(format)
    # The exporter class owns its configuration: it picks the settings its format reads out of this method's
    # union-of-every-format signature and drops the rest, so no dispatcher here has to know which is which.
    config = exporter_class.build_config(
        output_dir=Path(output_dir),
        output_name=output_name,
        variant_name=getattr(self, "size", None),
        backbone_only=backbone_only,
        dynamic_batch=dynamic_batch,
        verbose=verbose,
        notes=notes,
        opset_version=opset_version,
        backend=backend,
        soc=soc,
        fp16=fp16,
        coreml_precision=coreml_precision,
        coreml_neural_engine=coreml_neural_engine,
        coreai_precision=coreai_precision,
        openvino_precision=openvino_precision,
        quantization=quantization,
        calibration_data=calibration_data,
        max_images=max_images,
        batch_size=export_batch_size,
        max_batch_size=export_max_batch_size,
        trt_metadata=trt_metadata,
        trt_hardware_compatibility=trt_hardware_compatibility,
        trt_version_compatible=trt_version_compatible,
        trt_timing_cache=trt_timing_cache,
    )
    # Constructing the exporter validates the format's own settings (precision, quantization, notes, ...), then
    # warns about the ones it ignores.
    exporter = exporter_class(config)
    # The request holds up; now the host must too. A missing install is refused here rather than inside the
    # conversion, which is reached only after prepare_export_graph's full forward pass below. It follows every check
    # of the request above, so an invalid request is reported as one whether or not the format's dependency happens
    # to be installed: installing an extra would not help it. The configuration-dependent check comes last, once the
    # packages it may import are known to be there.
    exporter_class.check_dependencies()
    exporter.check_environment()
    logger.info(f"Exporting model to {format} format")

    device = self.model.device

    # Move the live model to CPU before deepcopying and keep it there during export. ``nn.Module.to(...)`` mutates
    # in place, so this frees GPU memory for the local export copy, ONNX tracing, TFLite conversion, and any
    # calibration tensors. The ``finally`` block restores the live model even if export or conversion raises.
    # Both moves go through ``_locked_move``: they rewrite the parameters a concurrent first ``predict()`` may be
    # moving to the accelerator at the same time. The span between them is not locked — a ``predict()`` arriving
    # mid-export can move the weights back onto the accelerator, costing this staging its freed memory.
    self.model.model = _locked_move(self.model.model, "cpu")
    model = deepcopy(self.model.model)
    model.to(device)
    try:
        os.makedirs(output_dir, exist_ok=True)
        from rfdetr.export._runtime.metadata import metadata_from_model

        try:
            metadata = metadata_from_model(
                self,
                format=format,
                shape=shape,
                batch_size=export_batch_size,
                dynamic_batch=dynamic_batch,
                backbone_only=backbone_only,
                backend=backend,
                max_batch_size=export_max_batch_size,
            )
        except ValidationError as error:
            logger.warning(
                "Could not capture inference metadata: %s. Export will proceed without inference metadata; "
                "load the artifact with RFDETRInference(metadata=<json>).",
                error,
                exc_info=True,
            )
            metadata = None

        graph = prepare_export_graph(
            model,
            self.model_config,
            shape=shape,
            device=device,
            infer_dir=infer_dir,
            batch_size=export_batch_size,
            dynamic_batch=dynamic_batch,
            backbone_only=backbone_only,
            metadata=metadata,
        )
        return exporter(graph)
    finally:
        self.model.model = _locked_move(self.model.model, device)

export_for_roboflow(output_dir)

Write a Roboflow upload bundle (weights.pt + class_names.txt) into output_dir.

This is the network-free core of :meth:deploy_to_roboflow: it serialises the model state and a sanitized copy of the training args into weights.pt, always embedding class_names so the bundle is self-contained, and writes the class labels to class_names.txt. The Roboflow SDK uses this format to adapt raw PyTorch-Lightning checkpoints into a deploy-ready bundle.

Parameters:

Name Type Description Default

output_dir

str | PathLike[str]

Directory into which weights.pt and class_names.txt are written. Created if it does not exist. Existing files are silently overwritten.

required

Raises:

Type Description
PermissionError

If the process lacks write access to output_dir or its parent directory.

OSError

On disk-full, invalid path, or other filesystem failure during directory creation, file write, or torch.save.

RuntimeError

If a legacy in-place optimization cleared the base weights.

Source code in src/rfdetr/detr.py
def export_for_roboflow(self, output_dir: str | os.PathLike[str]) -> None:
    """Write a Roboflow upload bundle (``weights.pt`` + ``class_names.txt``) into *output_dir*.

    This is the network-free core of :meth:`deploy_to_roboflow`: it serialises the model state and
    a sanitized copy of the training args into ``weights.pt``, always embedding ``class_names`` so
    the bundle is self-contained, and writes the class labels to ``class_names.txt``.  The Roboflow
    SDK uses this format to adapt raw PyTorch-Lightning checkpoints into a deploy-ready bundle.

    Args:
        output_dir: Directory into which ``weights.pt`` and ``class_names.txt`` are written.  Created
            if it does not exist.  Existing files are silently overwritten.

    Raises:
        PermissionError: If the process lacks write access to *output_dir* or its parent directory.
        OSError: On disk-full, invalid path, or other filesystem failure during directory creation,
            file write, or ``torch.save``.
        RuntimeError: If a legacy in-place optimization cleared the base weights.
    """
    if getattr(self, "_optimized_inplace", False) or self.model.model is None:
        raise RuntimeError(
            "Cannot export after legacy inplace optimization — "
            "the model has been cleared from memory. "
            "Call export_for_roboflow() before optimizing."
        )
    os.makedirs(output_dir, exist_ok=True)
    # Write class_names.txt so the Roboflow upload pipeline can discover
    # the class labels without relying on args.class_names in the checkpoint.
    class_names_path = os.path.join(output_dir, "class_names.txt")
    with open(class_names_path, "w", encoding="utf-8", newline="\n") as f:
        f.write("\n".join(self.class_names))

    # Serialize a sanitized copy so trained bundles do not expose the caller's
    # filesystem layout.  Keep self.model.args unchanged for runtime consumers.
    args = copy(self.model.args)
    args.dataset_dir = None
    args.output_dir = "output"
    args.resume = ""
    if not hasattr(args, "class_names") or args.class_names is None:
        args.class_names = self.class_names

    outpath = os.path.join(output_dir, "weights.pt")
    torch.save({"model": self.model.model.state_dict(), "args": args}, outpath)

from_checkpoint(path, *, trust_checkpoint=False, **kwargs) classmethod

Load an RF-DETR model from a training checkpoint, automatically inferring the model class.

The correct subclass is resolved in order of preference:

  1. model_name key in the checkpoint (written by the PTL training stack since v1.7.0).
  2. pretrain_weights field in the checkpoint's args entry (legacy fallback for older checkpoints).
  3. The filename of path itself, used as a last resort when pretrain_weights is absent or an unset-like sentinel value (empty string, "none", or "null"). Starter weights published by Roboflow store pretrain_weights="none" in their args; passing the canonical filename (e.g. rf-detr-small.pth) lets from_checkpoint infer the class automatically.

Both legacy argparse.Namespace checkpoints (produced by engine.py) and dict-style checkpoints (produced by the PTL training stack) are supported.

Parameters:

Name Type Description Default

path

str | PathLike[str]

Path to a checkpoint file (e.g. checkpoint_best_total.pth, or last.ckpt for the latest non-EMA weights of a run).

required

trust_checkpoint

bool

When True, fall back to weights_only=False (full pickle) if safe deserialization fails. Only set this for checkpoints from fully trusted sources; the default False keeps the safe loading path and raises if it cannot succeed. Applies to both the initial checkpoint read here and the constructor's own reload of the same file via :func:~rfdetr.models.weights.load_pretrain_weights.

False

**kwargs

Any

Additional keyword arguments forwarded to the model constructor (e.g. accept_platform_model_license=True for XLarge / 2XLarge models).

num_classes is resolved in this priority order:

  1. Explicit caller kwarg — always wins.
  2. Weight inference from class_embed.weight shape in the checkpoint (shape[0] - 1, since the head includes a background class). This overrides a stale model_config value written before fine-tuning changed the class count.
  3. saved_model_config["num_classes"] from the checkpoint's model_config entry — may be stale for older checkpoints.
  4. Legacy args["num_classes"] dict entry.
  5. Constructor default.

In cases 2–5 the field is not recorded as a user-set override, so :meth:train can still adapt the detection head to the training dataset's class count. Pass an explicit num_classes=N to pin the head and prevent adaptation.

The checkpoint's other model_config fields (resolution, num_select, dec_layers and the rest of the trained architecture) are restored the same way: an explicit caller kwarg always wins over the saved value. device and pretrain_weights are the two exceptions — path itself supplies the weights, and device is never restored, so the loading host's own default applies unless device= is passed, letting a GPU-trained checkpoint load on a CPU-only machine. When the checkpoint carries no model_config at all (best-total files written before it was strip-preserved), the lost fields fall back to class defaults and a warning names them, together with the unstripped sibling checkpoint or training_config.json to recover them from.

{}

Returns:

Type Description
RFDETR

An instance of the appropriate :class:RFDETR subclass loaded from the checkpoint.

Warning

By default this method attempts safe deserialization (weights_only=True). Pass trust_checkpoint=True only for checkpoints from fully trusted sources, as it enables full pickle deserialization which can execute arbitrary code.

Raises:

Type Description
FileNotFoundError

If path does not exist.

OSError

If path exists but cannot be read.

ValueError

If the model class cannot be inferred from model_name, pretrain_weights, or the checkpoint filename, or if path records no "args" at all — as last.ckpt and checkpoint_<epoch>.ckpt from rfdetr 1.11.0 and earlier are, and as :func:~rfdetr.training.checkpoint.convert_legacy_checkpoint output is.

ImportError

If the checkpoint was saved by a backbone_lora=True run and peft is not installed (pip install "rfdetr[lora]").

Examples:

model = RFDETR.from_checkpoint("checkpoint_best_total.pth") model = RFDETRSmall.from_checkpoint("checkpoint_best_total.pth")

Source code in src/rfdetr/detr.py
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
@classmethod
def from_checkpoint(cls, path: str | os.PathLike[str], *, trust_checkpoint: bool = False, **kwargs: Any) -> RFDETR:
    """Load an RF-DETR model from a training checkpoint, automatically inferring the model class.

    The correct subclass is resolved in order of preference:

    1. ``model_name`` key in the checkpoint (written by the PTL training
       stack since v1.7.0).
    2. ``pretrain_weights`` field in the checkpoint's ``args`` entry
       (legacy fallback for older checkpoints).
    3. The **filename** of *path* itself, used as a last resort when
       ``pretrain_weights`` is absent or an unset-like sentinel value
       (empty string, ``"none"``, or ``"null"``).  Starter weights
       published by Roboflow store ``pretrain_weights="none"`` in their
       ``args``; passing the canonical filename (e.g.
       ``rf-detr-small.pth``) lets ``from_checkpoint`` infer the class
       automatically.

    Both legacy ``argparse.Namespace`` checkpoints (produced by ``engine.py``) and dict-style checkpoints (produced
    by the PTL training stack) are supported.

    Args:
        path: Path to a checkpoint file (e.g. ``checkpoint_best_total.pth``, or ``last.ckpt`` for the latest
            non-EMA weights of a run).
        trust_checkpoint: When ``True``, fall back to ``weights_only=False``
            (full pickle) if safe deserialization fails.  Only set this for
            checkpoints from fully trusted sources; the default ``False``
            keeps the safe loading path and raises if it cannot succeed.
            Applies to both the initial checkpoint read here and the
            constructor's own reload of the same file via
            :func:`~rfdetr.models.weights.load_pretrain_weights`.
        **kwargs: Additional keyword arguments forwarded to the model
            constructor (e.g. ``accept_platform_model_license=True`` for XLarge / 2XLarge models).

            ``num_classes`` is resolved in this priority order:

            1. Explicit caller kwarg — always wins.
            2. Weight inference from ``class_embed.weight`` shape in the checkpoint
               (``shape[0] - 1``, since the head includes a background class). This
               overrides a stale ``model_config`` value written before fine-tuning
               changed the class count.
            3. ``saved_model_config["num_classes"]`` from the checkpoint's
               ``model_config`` entry — may be stale for older checkpoints.
            4. Legacy ``args["num_classes"]`` dict entry.
            5. Constructor default.

            In cases 2–5 the field is not recorded as a user-set override, so
            :meth:`train` can still adapt the detection head to the training
            dataset's class count.  Pass an explicit ``num_classes=N`` to pin
            the head and prevent adaptation.

            The checkpoint's other ``model_config`` fields (``resolution``,
            ``num_select``, ``dec_layers`` and the rest of the trained
            architecture) are restored the same way: an explicit caller kwarg
            always wins over the saved value.  ``device`` and
            ``pretrain_weights`` are the two exceptions — *path* itself supplies
            the weights, and ``device`` is never restored, so the loading host's
            own default applies unless ``device=`` is passed, letting a
            GPU-trained checkpoint load on a CPU-only machine.  When the checkpoint carries no
            ``model_config`` at all (best-total files written before it was
            strip-preserved), the lost fields fall back to class defaults and a
            warning names them, together with the unstripped sibling checkpoint
            or ``training_config.json`` to recover them from.

    Returns:
        An instance of the appropriate :class:`RFDETR` subclass loaded from the checkpoint.

    Warning:
        By default this method attempts safe deserialization
        (``weights_only=True``).  Pass ``trust_checkpoint=True`` only for
        checkpoints from fully trusted sources, as it enables full pickle
        deserialization which can execute arbitrary code.

    Raises:
        FileNotFoundError: If *path* does not exist.
        OSError: If *path* exists but cannot be read.
        ValueError: If the model class cannot be inferred from ``model_name``,
            ``pretrain_weights``, or the checkpoint filename, or if *path* records no ``"args"`` at all — as
            ``last.ckpt`` and ``checkpoint_<epoch>.ckpt`` from rfdetr 1.11.0 and earlier are, and as
            :func:`~rfdetr.training.checkpoint.convert_legacy_checkpoint` output is.
        ImportError: If the checkpoint was saved by a ``backbone_lora=True`` run and ``peft``
            is not installed (``pip install "rfdetr[lora]"``).

    Examples:
        model = RFDETR.from_checkpoint("checkpoint_best_total.pth")
        model = RFDETRSmall.from_checkpoint("checkpoint_best_total.pth")
    """
    # Local import breaks the variants → detr import cycle.
    import rfdetr.variants as rfdetr_variants

    _plus_available = False
    # Raised only if the checkpoint needs a plus class: a broken rfdetr_plus must not block core checkpoints.
    _plus_import_error: ImportError | None = None
    _plus_symbols: dict[str, type[RFDETR]] = {}
    # Why each plus symbol is missing; chained onto the upgrade hint below so the cause stays visible.
    _plus_symbol_errors: dict[str, ImportError] = {}
    _plus_entries: list[tuple[str, type[RFDETR]]] = []
    _plus_stem_entries: list[tuple[str, type[RFDETR]]] = []
    from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE

    if _IS_RFDETR_PLUS_AVAILABLE:
        try:
            import rfdetr.platform.models as platform_models

            for class_symbol in sorted(_PLUS_EXPORTS):
                try:
                    plus_obj = getattr(platform_models, class_symbol)
                except ImportError as ex:
                    # The installed rfdetr_plus predates this model; a checkpoint naming it is rejected below.
                    _plus_symbol_errors[class_symbol] = ex
                    continue
                _plus_symbols[class_symbol] = plus_obj
            _plus_entries = [
                (name, _plus_symbols[class_symbol])
                for name, class_symbol in _CHECKPOINT_PLUS_MODEL_MAP_ENTRIES
                if class_symbol in _plus_symbols
            ]
            _plus_stem_entries = [
                (name, _plus_symbols[class_symbol])
                for name, class_symbol in _CHECKPOINT_PLUS_STEM_ENTRIES
                if class_symbol in _plus_symbols
            ]
            _plus_available = True
        except ModuleNotFoundError as ex:
            if ex.name not in {"rfdetr_plus", "rfdetr_plus.models"}:
                _plus_import_error = ex
        except ImportError as ex:
            _plus_import_error = ex
        if _plus_import_error is not None:
            # Every from_checkpoint call retries the import; report the broken install once per process.
            logger.warning_once(
                "rfdetr_plus is installed but failed to import (%s); plus model checkpoints cannot be loaded.",
                _plus_import_error,
            )

    # Use the safe-load helper which tries weights_only=True first (with
    # legacy argparse.Namespace safe globals), falling back to full pickle
    # only when the caller explicitly passes trust_checkpoint=True.
    from rfdetr.utilities.io import _safe_torch_load

    ckpt: dict[str, Any] = _safe_torch_load(str(path), trust=trust_checkpoint)
    if "args" not in ckpt:
        raise ValueError(
            f"Checkpoint {os.fspath(path)!r} records no 'args' or 'model_config' — as last.ckpt and "
            "checkpoint_<epoch>.ckpt from rfdetr 1.11.0 and earlier are, and as convert_legacy_checkpoint "
            "output is (it stores the training args under 'hyper_parameters' instead) — so neither the model "
            "class nor its settings can be read from it. Load checkpoint_best_ema.pth from the same run instead "
            "(checkpoint_best_regular.pth if it trained without EMA), which records them since rfdetr 1.8.0. Or "
            "pass this file to the class it was trained with, with the architecture settings it was trained with, "
            f"e.g. RFDETRSmall(pretrain_weights={os.fspath(path)!r}, resolution=...); training_config.json lists "
            "them when the run wrote one. That route restores the weights but not the class names, which this "
            "file does not store."
        )
    args = ckpt["args"]

    _variant_name_to_class: dict[str, type[RFDETR]] = {
        getattr(variant_obj, "__name__", symbol): variant_obj
        for symbol in dir(rfdetr_variants)
        if symbol.startswith("RFDETR")
        for variant_obj in [getattr(rfdetr_variants, symbol)]
    }
    _variant_symbols: dict[str, type[RFDETR]] = {
        class_symbol: _variant_name_to_class[class_symbol] for class_symbol in _CHECKPOINT_MODEL_NAME_CLASS_SYMBOLS
    }
    # Build in three explicit segments: seg-* entries, then plus-model entries
    # (xlarge/2xlarge), then base entries — order determines lookup priority.
    _seg_map: list[tuple[str, type[RFDETR]]] = [
        (name, _variant_symbols[class_symbol])
        for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
        if name.startswith("seg-")
    ]
    _keypoint_map: list[tuple[str, type[RFDETR]]] = [
        (name, _variant_symbols[class_symbol])
        for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
        if "keypoint" in name
    ]
    _base_map: list[tuple[str, type[RFDETR]]] = [
        (name, _variant_symbols[class_symbol])
        for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
        if not name.startswith("seg-") and "keypoint" not in name
    ]
    _model_map: list[tuple[str, type[RFDETR]]] = _seg_map + _keypoint_map + _plus_entries + _base_map

    # New checkpoints store model_name directly — use it when available.
    _name_map: dict[str, type[RFDETR]] = dict(_variant_symbols)
    # Plus-model classes are resolved only when rfdetr_plus is installed.
    if _plus_available:
        _name_map.update(_plus_symbols)
    # RFDETRLargeDeprecated is excluded from _CHECKPOINT_MODEL_NAME_CLASS_SYMBOLS
    # (so forward name-map lookups always reach RFDETRLarge), but it must be
    # in _name_map so that checkpoints carrying model_name="RFDETRLargeDeprecated"
    # are reloaded with the correct class instead of falling through to the
    # substring matcher (which would wrongly pick RFDETRLarge and fail with a
    # pydantic literal_error on encoder / projector_scale fields).
    # Note: RFDETRLargeDeprecated is a _DeprecatedProxy (pyDeprecate) and has no
    # __name__; look it up by the string key it was registered under, then inject
    # into _name_map so checkpoint matching resolves directly without the substring
    # fallback.  Tests assert this mapping exists (see tests/inference/test_from_checkpoint.py).
    _large_deprecated_cls = _variant_name_to_class.get("RFDETRLargeDeprecated")
    if _large_deprecated_cls is not None:
        _name_map["RFDETRLargeDeprecated"] = _large_deprecated_cls
    saved_model_name = ckpt.get("model_name")
    model_cls: type[RFDETR] | None = None
    if isinstance(saved_model_name, str):
        normalized_name = saved_model_name.strip()
        if normalized_name:
            model_cls = _name_map.get(normalized_name)
    else:
        normalized_name = ""

    # Fall back to pretrain_weights (legacy) or, when unset-like, the checkpoint filename.
    if isinstance(args, dict):
        weights_name = str(args.get("pretrain_weights", "")).strip().lower()
    else:
        weights_name = str(getattr(args, "pretrain_weights", "")).strip().lower()
    # The sentinel set {"", "none", "null"} covers unset-like checkpoint values:
    #   ""     — pretrain_weights key absent entirely
    #   "none" — checkpoint value was None or the literal string "none";
    #            after str(...).strip().lower() both normalize to the same sentinel.
    #            This is NOT an intentional "no pretraining" flag (see
    #            test_pretrain_weights_none_warns, which operates at the config
    #            level, not the checkpoint level)
    #   "null" — checkpoint stored the literal string "null" (for example from a
    #            YAML-originated value), which is also treated as unset-like here
    _filename_fallback = False
    if weights_name in {"", "none", "null"}:
        weights_name = os.path.basename(os.fspath(path)).lower()
        _filename_fallback = True
    # A directory such as "rf-detr-pico/" names no model. ntpath splits on both "/" and "\", so a Windows-style
    # path recorded in the checkpoint (e.g. "c:\models\rf-detr-pico\rf-detr-nano.pth") also drops its directories.
    weights_file = ntpath.basename(weights_name)

    if model_cls is None:
        # Guard: plus-only checkpoints should raise an actionable install error
        # when rfdetr_plus is missing, regardless of whether class inference
        # relies on model_name (new format) or pretrain_weights (legacy format).
        plus_by_model_name = normalized_name in _PLUS_EXPORTS
        plus_by_weights_name = (
            "xlarge" in weights_name and "seg-" not in weights_name and "keypoint-preview" not in weights_name
        ) or any(name in weights_file for name, _ in _CHECKPOINT_PLUS_STEM_ENTRIES)
        if not _plus_available and (plus_by_model_name or plus_by_weights_name):
            from rfdetr.platform import _INSTALL_MSG

            reason = (
                f", which is installed but failed to import: {_plus_import_error}"
                if _plus_import_error
                else ". " + _INSTALL_MSG.format(name="platform model downloads")
            )
            raise ImportError(
                f"Checkpoint model_name={saved_model_name!r}, pretrain_weights={weights_name!r} requires the "
                f"rfdetr_plus package{reason}"
            ) from _plus_import_error
        # An installed rfdetr_plus that predates the checkpoint's model (by name or release stem): never fall back
        # to other names, which could resolve a core class (e.g. "small" in "rf-detr-pico-small-ft.pth").
        missing_symbol = normalized_name if plus_by_model_name else None
        for stem, symbol in _CHECKPOINT_PLUS_STEM_ENTRIES:
            if missing_symbol is None and stem in weights_file and symbol not in _plus_symbols:
                missing_symbol = symbol
        if missing_symbol is not None:
            from rfdetr.platform.models import _UPGRADE_MSG

            raise ImportError(
                f"Checkpoint model_name={saved_model_name!r}, pretrain_weights={weights_name!r}: "
                + _UPGRADE_MSG.format(name=missing_symbol)
            ) from _plus_symbol_errors.get(missing_symbol)

        model_cls = next((klass for name, klass in _plus_stem_entries if name in weights_file), None)
        if model_cls is None:
            for name, klass in _model_map:
                if name in weights_name:
                    model_cls = klass
                    break

        if _filename_fallback and model_cls is not None:
            logger.info(
                "pretrain_weights unset in checkpoint %r; inferred model class %s from filename %r",
                path,
                getattr(model_cls, "__name__", repr(model_cls)),
                weights_name,
            )

    if model_cls is None:
        raise ValueError(
            f"Could not infer model class from checkpoint at {path!r} "
            f"(model_name={saved_model_name!r}, pretrain_weights={weights_name!r}). "
            f"Please instantiate the model class directly."
        )

    if isinstance(args, dict):
        num_classes: int | None = args.get("num_classes")
    else:
        num_classes = getattr(args, "num_classes", None)

    constructor_kwargs: dict[str, Any] = {}
    checkpoint_config_keys: set[str] = set()  # keys injected from checkpoint, not from caller

    # Resolve model config field set once — used for both saved_model_config parsing and
    # weight-based schema inference guards (BaseConfig has extra="forbid"; unknown fields raise).
    _model_config_class = getattr(model_cls, "_model_config_class", None)
    _mc_fields: dict[str, Any] = {}
    _mc_model_fields = getattr(_model_config_class, "model_fields", None)
    if isinstance(_mc_model_fields, dict):
        _mc_fields = _mc_model_fields
    else:
        _mc_legacy = getattr(_model_config_class, "__fields__", None)
        if isinstance(_mc_legacy, dict):
            _mc_fields = _mc_legacy

    saved_model_config = ckpt.get("model_config")
    if isinstance(saved_model_config, dict):
        for key, value in saved_model_config.items():
            # device records the training host (e.g. "cuda"); the loading host keeps its own default.
            if key in ("pretrain_weights", "device"):
                continue
            if not _mc_fields or key in _mc_fields:
                constructor_kwargs[key] = value
                checkpoint_config_keys.add(key)
    elif "rfdetr_version" in ckpt or "best_total_source" in ckpt:
        # checkpoint_best_total.pth files written before strip_checkpoint kept model_config have lost it; the
        # unstripped file they were copied from, named by best_total_source, still has it. These settings then
        # fall back to class defaults without an error (a changed num_queries or group_detr fails loudly on the
        # weight shapes instead), so keep warning until the caller has passed all of them.
        # Both keys discriminate, because neither is present on its own in every affected file: rfdetr_version
        # reaches back to 1.7.0 but is omitted when get_version() cannot resolve a version (editable install
        # without package metadata), while best_total_source covers those but only exists since 1.9.0.
        silent_fields = ["resolution", "num_select", "dec_layers"]
        segmentation_field = _mc_fields.get("segmentation_head")
        if kwargs.get("segmentation_head", getattr(segmentation_field, "default", False)) is True:
            silent_fields.append("mask_downsample_ratio")
        missing = [name for name in silent_fields if name not in kwargs]
        if missing:
            # A best-total file often travels alone (the Roboflow SDK upload path sends only this one), so the
            # sibling it was copied from is not always there to point at. Only the two source names the training
            # stack writes become a filename; anything else is reported as-is so an anomaly stays visible.
            best_total_source = ckpt.get("best_total_source")
            sibling = (
                Path(path).with_name(f"checkpoint_best_{best_total_source}.pth")
                if best_total_source in ("ema", "regular")
                else None
            )
            if sibling is not None and sibling.exists():
                remedy = (
                    f"Load {sibling.name} from the same output directory instead, or pass the training values "
                    "to from_checkpoint(); training_config.json in that directory lists them."
                )
            elif sibling is not None:
                remedy = (
                    f"{sibling.name} is not beside it, so pass the values above to from_checkpoint() "
                    "explicitly; training_config.json from the original training run lists them."
                )
            else:
                remedy = (
                    f"Its best_total_source is {best_total_source!r}, so check for an unstripped "
                    "checkpoint_best_*.pth beside it, or pass the values above to from_checkpoint() "
                    "explicitly; training_config.json from the original training run lists them."
                )
            written_by = f" (written by rfdetr {ckpt['rfdetr_version']})" if "rfdetr_version" in ckpt else ""
            logger.warning(
                "Checkpoint %r%s has no model_config, which checkpoint_best_total.pth files lost when "
                "stripped, so these settings fall back to %s defaults: %s. %s",
                str(path),
                written_by,
                getattr(model_cls, "__name__", repr(model_cls)),
                ", ".join(missing),
                remedy,
            )

    if num_classes is not None and "num_classes" not in kwargs:
        constructor_kwargs["num_classes"] = num_classes
        checkpoint_config_keys.add("num_classes")

    # Infer schema-critical fields from checkpoint weights — these are authoritative when
    # ``model_config`` is absent or stale (saved before ``model_config`` persistence was added,
    # or saved with default values before fine-tuning changed the trained schema).
    # User-supplied ``kwargs`` take precedence and are applied in the ``update`` call below.
    _ckpt_weights: dict[str, Any] = ckpt.get("model") or {}
    if not _ckpt_weights and "state_dict" in ckpt:
        _pfx = "model."
        _ckpt_weights = {}
        for k, v in ckpt["state_dict"].items():
            if k.startswith(_pfx):
                key = k[len(_pfx) :]
                # Strip optional torch.compile() wrapper prefix
                if key.startswith("_orig_mod."):
                    key = key[len("_orig_mod.") :]
                _ckpt_weights[key] = v
    if _ckpt_weights:
        # num_keypoints_per_class — inferred from _kp_active_mask (shape [num_classes, max_kp]).
        # Reflects what the model actually learned; saved model_config may carry the COCO default
        # [0, 17] even after fine-tuning on a different keypoint schema.
        if "num_keypoints_per_class" not in kwargs and (not _mc_fields or "num_keypoints_per_class" in _mc_fields):
            _kp_mask = _ckpt_weights.get("_kp_active_mask")
            if isinstance(_kp_mask, torch.Tensor) and _kp_mask.ndim == 2:
                _inferred_kp = [int(n) for n in _kp_mask.sum(dim=1).tolist()]
                _current_kp = constructor_kwargs.get("num_keypoints_per_class")
                if _inferred_kp != _current_kp:
                    logger.debug(
                        "from_checkpoint: overriding num_keypoints_per_class %s → %s "
                        "(inferred from _kp_active_mask; saved model_config may be stale).",
                        _current_kp,
                        _inferred_kp,
                    )
                constructor_kwargs["num_keypoints_per_class"] = _inferred_kp
                checkpoint_config_keys.add("num_keypoints_per_class")
        # num_classes — inferred from class_embed.weight shape.
        # The head shape is ground truth for what num_classes the checkpoint uses.
        if "num_classes" not in kwargs:
            _ce_weight = _ckpt_weights.get("class_embed.weight")
            if isinstance(_ce_weight, torch.Tensor) and _ce_weight.ndim == 2:
                _inferred_nc = _ce_weight.shape[0] - 1  # shape[0] = num_classes + 1 (background)
                _current_nc = constructor_kwargs.get("num_classes")
                if _inferred_nc != _current_nc:
                    logger.debug(
                        "from_checkpoint: overriding num_classes %s → %s "
                        "(inferred from class_embed.weight; saved model_config may be stale).",
                        _current_nc,
                        _inferred_nc,
                    )
                constructor_kwargs["num_classes"] = _inferred_nc
                checkpoint_config_keys.add("num_classes")

    constructor_kwargs.update(kwargs)
    # pretrain_weights is placed after **kwargs so it always wins even if
    # a caller accidentally passes pretrain_weights inside kwargs.
    # Fully resolved, because the config sends a bare filename to the model cache, which is not the file read above.
    constructor_kwargs["pretrain_weights"] = os.path.realpath(path)
    # Model construction reloads this same file via load_pretrain_weights(); without this,
    # trust_checkpoint=True would bypass the safe-load only for the metadata read above and
    # then fail identically when the constructor re-reads pretrain_weights.
    constructor_kwargs["trust_checkpoint"] = trust_checkpoint

    # Fields injected from the checkpoint but not supplied by the caller must not be
    # treated as explicit user overrides in Pydantic's model_fields_set.  Downstream
    # alignment guards (e.g. _align_num_classes_from_dataset,
    # _align_keypoint_schema_from_dataset, load_pretrain_weights) all read
    # model_fields_set to decide whether to adapt model internals to the training
    # dataset — leaving checkpoint-derived fields marked as user-set breaks them.
    checkpoint_derived_keys = checkpoint_config_keys - set(kwargs)

    model = model_cls(**constructor_kwargs)
    # The instance now carries checkpoint (trained) weights; flag it so a later
    # train() call warns that it will restart from pretrain_weights, not continue.
    model._has_been_trained = True

    if checkpoint_derived_keys:
        loaded_config = getattr(model, "model_config", None)
        # model_fields_set is the public API and returns the live backing set
        # in Pydantic v2; fall back to the private attribute only if that changes.
        fields_set = getattr(loaded_config, "model_fields_set", None)
        if fields_set is None:
            fields_set = getattr(loaded_config, "__pydantic_fields_set__", None)
        if fields_set is not None:
            fields_set.difference_update(checkpoint_derived_keys)
        # Verify num_classes specifically — if Pydantic ever returns a snapshot instead
        # of the live backing set, this assertion will catch the silent regression before
        # it causes a training-time head-adaptation failure.
        if "num_classes" in checkpoint_derived_keys:
            assert "num_classes" not in getattr(loaded_config, "model_fields_set", set()), (
                "num_classes still in model_fields_set after checkpoint load; "
                "Pydantic may return a snapshot rather than the live backing set — "
                "switch to model_construct(_fields_set=...) for Pydantic v3 compatibility."
            )

    return model

get_model(config, *, trust_checkpoint=False)

Retrieve a model context from the provided architecture configuration.

Parameters:

Name Type Description Default

config

ModelConfig

Architecture configuration.

required

trust_checkpoint

bool

Forwarded to :func:~rfdetr.inference._build_model_context — set True only when config.pretrain_weights is a checkpoint the caller explicitly trusts (mirrors RFDETR.from_checkpoint(..., trust_checkpoint=True)).

False

Returns:

Type Description
ModelContext

ModelContext with model, postprocess, device, resolution, args, and class_names attributes.

Source code in src/rfdetr/detr.py
def get_model(self, config: ModelConfig, *, trust_checkpoint: bool = False) -> ModelContext:
    """Retrieve a model context from the provided architecture configuration.

    Args:
        config: Architecture configuration.
        trust_checkpoint: Forwarded to :func:`~rfdetr.inference._build_model_context` — set
            ``True`` only when ``config.pretrain_weights`` is a checkpoint the caller
            explicitly trusts (mirrors ``RFDETR.from_checkpoint(..., trust_checkpoint=True)``).

    Returns:
        ModelContext with model, postprocess, device, resolution, args, and class_names attributes.
    """
    return _build_model_context(config, trust_checkpoint=trust_checkpoint)

get_model_config(**kwargs)

Retrieve the configuration parameters used by the model.

Source code in src/rfdetr/detr.py
def get_model_config(self, **kwargs: Any) -> ModelConfig:
    """Retrieve the configuration parameters used by the model."""
    return self._model_config_class(**kwargs)

get_train_config(**kwargs)

Retrieve the configuration parameters that will be used for training.

Source code in src/rfdetr/detr.py
def get_train_config(self, **kwargs: Any) -> TrainConfig:
    """Retrieve the configuration parameters that will be used for training."""
    return self._train_config_class(**kwargs)

inference(compile=True, batch_size=1, dtype=torch.float32, *, compile_backend=None)

Create an independent, optimized predictor from the current base weights.

The returned predictor owns a snapshot of the weights, configuration, class names, normalization and postprocessor. Creating or using it leaves this model's weights, device, dtype, training mode and optimization state unchanged. Later training or edits to this model do not affect the snapshot. Do not mutate or train the source concurrently with snapshot creation. An existing legacy optimized runtime is not copied; the current base weights are the source of truth.

On MPS, the snapshot samples deformable attention with the fused F.grid_sample and, for TorchScript and compile=False predictors, folds each DINOv2 LayerScale into the Linear layer before it. At float32 its predictions match the source model's to float rounding, not bit for bit; in half precision they can differ by more, while staying at least as close to float32 as the source's.

Capture the return value and call predictor.predict(images). Calling model.predict(images) continues to use the source model. Each factory call allocates its own weight copy; release a predictor to release its runtime, or release the source model when only inference is needed. The former inplace option is removed.

Parameters:

Name Type Description Default

compile

bool

Compile the snapshot for fixed batch size and resolution. Set to False to use its eager export-mode forward, with a fixed resolution and flexible batch size.

True

batch_size

int

Positive batch size for compilation. Ignored when compile=False.

1

dtype

dtype | str

Floating-point dtype for the snapshot, as a torch.dtype or its string name.

float32

compile_backend

Literal['torchscript', 'cudagraph', 'inductor'] | None

None picks "cudagraph" on CUDA, "inductor" on MPS, when the installed PyTorch can compile for it, and "torchscript" everywhere else; compile=False ignores it. "torchscript" traces the snapshot. "cudagraph" additionally freezes the trace and captures its fixed-shape CUDA work for replay, cloning graph-owned outputs before returning them. It is CUDA-only and retains a graph-private memory pool. "inductor" uses torch.compile, with mode="reduce-overhead" on CUDA and the default mode elsewhere, and warms up the result before returning. On CUDA, this includes a second call and synchronization for CUDA Graph Trees. Compilation support and setup cost depend on the installed PyTorch version and device. TorchScript freezing does not guarantee bit-identical numerics. Create CUDA graph predictors before serving threads start: capture synchronizes the device and empties process-wide CUDA and pinned-host caches.

None

Returns:

Type Description
RFDETRInference

An inference-only facade owning the optimized snapshot.

Raises:

Type Description
TypeError

If dtype is neither a floating-point dtype nor the name of a torch dtype.

ValueError

If the dtype is non-floating, the backend is unknown, the compiled batch size is invalid, or CUDA graph compilation targets a non-CUDA device.

RuntimeError

If the source has no base weights, or compilation/capture fails. A failed CUDA graph capture can leave CUDA random-number state unusable; restart the process before trying another backend.

OutOfMemoryError

If copying or compiling the snapshot exhausts device memory.

Examples:

Requires model weights; the predictor can be reused for multiple images.

>>> from rfdetr import RFDETRSmall
>>> model = RFDETRSmall()
>>> predictor = model.inference(compile=False)
>>> detections = predictor.predict("image.jpg")
Source code in src/rfdetr/detr.py
def inference(
    self,
    compile: bool = True,
    batch_size: int = 1,
    dtype: torch.dtype | str = torch.float32,
    *,
    compile_backend: Literal["torchscript", "cudagraph", "inductor"] | None = None,
) -> RFDETRInference:
    """Create an independent, optimized predictor from the current base weights.

    The returned predictor owns a snapshot of the weights, configuration, class names, normalization and
    postprocessor. Creating or using it leaves this model's weights, device, dtype, training mode and optimization
    state unchanged. Later training or edits to this model do not affect the snapshot. Do not mutate or train the
    source concurrently with snapshot creation. An existing legacy optimized runtime is not copied; the current
    base weights are the source of truth.

    On MPS, the snapshot samples deformable attention with the fused ``F.grid_sample`` and, for TorchScript and
    ``compile=False`` predictors, folds each DINOv2 LayerScale into the Linear layer before it. At float32 its
    predictions match the source model's to float rounding, not bit for bit; in half precision they can differ by
    more, while staying at least as close to float32 as the source's.

    Capture the return value and call ``predictor.predict(images)``. Calling ``model.predict(images)`` continues
    to use the source model. Each factory call allocates its own weight copy; release a predictor to release its
    runtime, or release the source model when only inference is needed. The former ``inplace`` option is removed.

    Args:
        compile: Compile the snapshot for fixed batch size and resolution. Set to ``False`` to use its eager
            export-mode forward, with a fixed resolution and flexible batch size.
        batch_size: Positive batch size for compilation. Ignored when ``compile=False``.
        dtype: Floating-point dtype for the snapshot, as a ``torch.dtype`` or its string name.
        compile_backend: ``None`` picks ``"cudagraph"`` on CUDA, ``"inductor"`` on MPS, when the installed
            PyTorch can compile for it, and ``"torchscript"`` everywhere else; ``compile=False`` ignores it.
            ``"torchscript"`` traces the snapshot. ``"cudagraph"`` additionally
            freezes the trace and captures its fixed-shape CUDA work for replay, cloning graph-owned outputs before
            returning them. It is CUDA-only and retains a graph-private memory pool. ``"inductor"`` uses
            ``torch.compile``, with ``mode="reduce-overhead"`` on CUDA and the default mode elsewhere, and warms up
            the result before returning. On CUDA, this includes a second call and synchronization for CUDA Graph
            Trees. Compilation support and setup cost
            depend on the installed PyTorch version and device. TorchScript freezing does not guarantee
            bit-identical numerics. Create CUDA graph predictors before serving threads start: capture synchronizes
            the device and empties process-wide CUDA and pinned-host caches.

    Returns:
        An inference-only facade owning the optimized snapshot.

    Raises:
        TypeError: If ``dtype`` is neither a floating-point dtype nor the name of a torch dtype.
        ValueError: If the dtype is non-floating, the backend is unknown, the compiled batch size is invalid,
            or CUDA graph compilation targets a non-CUDA device.
        RuntimeError: If the source has no base weights, or compilation/capture fails. A failed CUDA graph
            capture can leave CUDA random-number state unusable; restart the process before trying another backend.
        torch.cuda.OutOfMemoryError: If copying or compiling the snapshot exhausts device memory.

    Examples:
        Requires model weights; the predictor can be reused for multiple images.

        >>> from rfdetr import RFDETRSmall  # doctest: +SKIP
        >>> model = RFDETRSmall()  # doctest: +SKIP
        >>> predictor = model.inference(compile=False)  # doctest: +SKIP
        >>> detections = predictor.predict("image.jpg")  # doctest: +SKIP
    """
    if isinstance(dtype, str):
        try:
            dtype = getattr(torch, dtype)
        except AttributeError:
            raise TypeError(f"dtype must be a torch.dtype or a string name of a dtype, got {dtype!r}") from None
    if not isinstance(dtype, torch.dtype):
        raise TypeError(f"dtype must be a torch.dtype or a string name of a dtype, got {type(dtype)!r}")
    if not dtype.is_floating_point:
        raise ValueError(f"dtype must be a floating-point torch.dtype or string name of one, got {dtype}")
    if compile_backend not in (None, "torchscript", "cudagraph", "inductor"):
        raise ValueError(
            f"compile_backend must be 'torchscript', 'cudagraph', 'inductor' or None, got {compile_backend!r}"
        )
    if compile and (isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size < 1):
        raise ValueError("batch_size must be a positive integer when compile=True.")
    device = torch.device(self.model.device)
    if compile_backend is None:
        compile_backend = _default_compile_backend(device)
    if compile and compile_backend == "cudagraph" and device.type != "cuda":
        raise ValueError("compile_backend='cudagraph' requires a CUDA device.")
    source_model = self.model.model
    if source_model is None:
        raise RuntimeError("Cannot create a predictor: the base model has been cleared. Reload the source model.")
    if device.type == "cuda" and device.index is None:
        device = torch.device("cuda", torch.cuda.current_device())

    # Copy only prediction state, never training callbacks or a trainer which can retain the source model.
    # Exclude both weight-bearing entries while copying metadata, then clone the base module exactly once.
    snapshot = object.__new__(RFDETR)
    cuda_ctx = torch.cuda.device(device) if device.type == "cuda" else contextlib.nullcontext()
    with cuda_ctx:
        with _device_move_lock(source_model), torch.inference_mode(False):
            snapshot.model_config, snapshot.model, snapshot.means, snapshot.stds = deepcopy(
                (self.model_config, self.model, self.means, self.stds),
                {id(source_model): None, id(self.model.inference_model): None},
            )
            inference_model: Any = deepcopy(source_model)
        snapshot.model.device = device
        inference_model = inference_model.to(device=device)
        inference_model.eval()
        inference_model.export()
        # Cast before the MPS rewrites so LayerScale folds in the snapshot's dtype: folding a float16 source and
        # widening afterwards would keep products that underflowed in float16.
        inference_model = inference_model.to(dtype=dtype)
        if device.type == "mps":
            # The snapshot never runs backward.
            _rewrite_for_mps_prediction(
                inference_model, fold_layer_scales=not (compile and compile_backend == "inductor")
            )
        if compile:
            dummy_input = torch.randn(
                batch_size,
                snapshot.model_config.num_channels,
                snapshot.model.resolution,
                snapshot.model.resolution,
                device=device,
                dtype=dtype,
            )
            if compile_backend in ("torchscript", "cudagraph"):
                if device.type != "cpu" and getattr(inference_model, "segmentation_head", None) is not None:
                    # The owned prediction snapshot never runs backward. Its native upsample avoids the
                    # deterministic-backward decomposition that TorchScript cannot trace on accelerators.
                    inference_model.segmentation_head._native_prediction_upsample = True
                inference_model = _trace_for_prediction(inference_model, dummy_input)
                if compile_backend == "cudagraph":
                    # The existing CUDA parity/latency measurements use a frozen trace; retain that capture path.
                    inference_model = torch.jit.freeze(inference_model)
                    inference_model = _CUDAGraphInferenceModel(inference_model, dummy_input, device)
            else:
                # `reduce-overhead` only adds CUDA graph trees; elsewhere it records nothing and warns on compile.
                inference_model = torch.compile(
                    inference_model, mode="reduce-overhead" if device.type == "cuda" else "default"
                )
                with torch.inference_mode():
                    # `predict()` passes inference tensors, and dynamo guards on that: a warm-up input built
                    # outside inference mode makes the first prediction compile again.
                    warmup_input = dummy_input.clone()
                    inference_model(warmup_input)
                    if device.type == "cuda":
                        # CUDA Graph Trees reserve memory on the first call and record on the second.
                        inference_model(warmup_input)
                        torch.cuda.synchronize(device)

    snapshot.model.inference_model = inference_model
    snapshot._is_optimized_for_inference = True
    snapshot._has_warned_about_not_being_optimized_for_inference = False
    snapshot._optimized_has_been_compiled = compile
    snapshot._optimized_batch_size = batch_size if compile else None
    snapshot._optimized_resolution = snapshot.model.resolution
    snapshot._optimized_dtype = dtype
    snapshot._optimized_inplace = False
    return RFDETRInference._from_owned_model(snapshot)

maybe_download_pretrain_weights()

Download pre-trained weights if they are not already downloaded.

Bare filenames (no directory component, e.g. rf-detr-base.pth) are resolved to the model cache directory — set the RF_HOME environment variable to override the location (default: ~/.roboflow/models). Resolution happens in ModelConfig.expand_path for explicitly-provided values, and here as a fallback for field defaults (which Pydantic does not validate by default).

Paths that already contain a directory component are used as-is; the parent directory is created if it does not yet exist.

Source code in src/rfdetr/detr.py
def maybe_download_pretrain_weights(self) -> None:
    """Download pre-trained weights if they are not already downloaded.

    Bare filenames (no directory component, e.g. ``rf-detr-base.pth``) are resolved to the model cache directory —
    set the ``RF_HOME`` environment variable to override the location (default: ``~/.roboflow/models``). Resolution
    happens in ``ModelConfig.expand_path`` for explicitly-provided values, and here as a fallback for field defaults
    (which Pydantic does not validate by default).

    Paths that already contain a directory component are used as-is; the parent directory is created if it does not
    yet exist.
    """
    if self.model_config.pretrain_weights is None:
        return
    pretrain_weights = str(self.model_config.pretrain_weights)
    if not os.path.dirname(pretrain_weights):
        # Field default was not processed by expand_path — resolve to cache dir.
        cache_dir = get_model_cache_dir()
        os.makedirs(cache_dir, exist_ok=True)
        pretrain_weights = os.path.join(cache_dir, pretrain_weights)
    else:
        os.makedirs(os.path.dirname(pretrain_weights), exist_ok=True)
    self.model_config.pretrain_weights = pretrain_weights
    download_pretrain_weights(pretrain_weights)

predict(images, threshold=0.5, shape=None, patch_size=None, include_source_image=True, *, antialias=False, **kwargs)

Performs model inference on the input images.

This method accepts a single image or an iterable of images (string or path-like file path, image URL string, PIL Image, NumPy array, or torch.Tensor) and runs them as one batch. The images should be in RGB channel order. A torch.Tensor must already be normalized to values in the [0, 1] range and have the shape (C, H, W); a (B, C, H, W) tensor is split into B images of one batch. Videos, directories, globs, manifests, webcams and streams are not images: predict them through RFDETRInference(model).predict(...).

Parameters:

Name Type Description Default

images

_ImageInput | Iterable[_ImageInput]

A single image or an iterable of images to process. Images can be provided as string or path-like file paths, URL strings, PIL Images, NumPy arrays, or torch.Tensors. Iterables other than lists and tuples are read into a list first. PIL and file inputs convert to grayscale for one-channel models and RGB for three-channel models. HTTP(S) URL strings are downloaded with requests using a 30-second connect and read timeout, which does not bound the total download time. Redirects are followed, and the whole response body is read into memory with no size limit or host restriction. Pass only URLs you trust; for untrusted URLs, download and validate the image yourself and pass the decoded image instead.

required

threshold

_ThresholdInput

The confidence score a detection must exceed to be kept. A single float applies to every class. A :class:~rfdetr.ClassThresholds, such as ClassThresholds({"person": 0.35}, default=0.5), applies each class's own threshold to its detections, so that a class the model finds less reliably can be held to a different bar. It leaves out every detection that belongs to no entry of class_names: the no-object slot of a fine-tuned model, the background slot of a legacy keypoint checkpoint, and the unused category IDs of a pretrained COCO model. A threshold of float("inf") drops every detection of that class.

0.5

shape

tuple[int, int] | None

Resize to (height, width); defaults to the model resolution. Both dimensions must be positive integers divisible by patch_size * num_windows. Use the model's export resolution when applicable.

None

patch_size

int | None

Backbone patch size for shape validation. Defaults to model_config.patch_size.

None

include_source_image

bool

Whether to attach the original image to the returned prediction. Detection and segmentation outputs use detections.metadata["source_image"]. Keypoint outputs use per-object key_points.data["source_image"] because Supervision KeyPoints currently has no collection-level metadata field. Defaults to True. Set to False to reduce memory use when source images are not needed.

True

antialias

bool

Whether to use antialiasing during the inference resize. Match it to the resize the checkpoint was trained with. Keep this False for checkpoints trained with the antialias-free Albumentations pipeline, which is what the default CPU augmentation backend uses when rfdetr[augment] is installed. Set it to True for checkpoints trained with torchvision resizing: the Kornia/GPU augmentation backend, the CPU backend without rfdetr[augment], older RF-DETR releases, or antialiased platform preprocessing. Whether rfdetr[augment] was installed at training time therefore decides which setting a default-trained checkpoint expects. RFDETRInference uses the artifact metadata's resize_antialias by default for exported models; callers can override it with an explicit antialias argument. The lower-level rfdetr.export runtime preprocessing helpers always resize with antialias=False. Defaults to False to preserve the existing inference behavior.

False

**kwargs

Any

Unknown keyword arguments are ignored with a FutureWarning. Passing them is deprecated since v1.12.0 and will raise TypeError in v2.0.0. The source options of :meth:RFDETRInference.predict (stream, batch, vid_stride, stream_buffer) raise TypeError now.

{}

Returns:

Type Description
_PredictionOutput

One result for a single image, or a list with one result per image for an iterable or a BCHW tensor.

_PredictionOutput

Detection and segmentation return :class:~supervision.Detections; keypoints return

_PredictionOutput

class:~supervision.KeyPoints with xy coordinates.

_PredictionOutput

key_points.detection_confidence is the per-object score used by threshold. For keypoint models this

_PredictionOutput

is the postprocessed detection score and, by default, includes normalized keypoint uncertainty fusion

_PredictionOutput

controlled by model_config.postprocess_trace_alpha. key_points.keypoint_confidence is separate: it

_PredictionOutput

is a (num_detections, num_keypoints) array of per-keypoint findability scores decoded from the keypoint

_PredictionOutput

head, not a repeated copy of the detection score. When RF-DETR emits keypoint precision parameters,

_PredictionOutput

key_points.data["covariance"] stores per-keypoint pixel-space covariance matrices with shape

_PredictionOutput

(num_detections, num_keypoints, 2, 2). key_points.data["xyxy"] stores the corresponding detection

_PredictionOutput

boxes as a (num_detections, 4) array in the same row order as key_points.xy because Supervision

_PredictionOutput

KeyPoints does not have a native bounding-box field. The data dict also contains class_name and

_PredictionOutput

source_shape as per-object arrays. When include_source_image=True for keypoint models,

_PredictionOutput

source_image is stored as per-object data until Supervision exposes collection-level metadata for

_PredictionOutput

KeyPoints.

Note

For Detections outputs, source_image moved from detections.data to detections.metadata. Update detection callers reading detections.data["source_image"] to use detections.metadata["source_image"].

Note

class_name mapping uses one of three modes depending on the checkpoint. For pretrained COCO checkpoints (detected when model.args.num_classes > len(class_names) and class_names matches COCO_CLASS_NAMES), raw COCO category IDs (1–90, sparse) are looked up by category ID rather than by position — so class_id=18 yields "dog", not class_names[18]. For fine-tuned detection and segmentation models and active-first keypoint models, class_id is a 0-based index into class_names. In the one-class preview keypoint setup, that means class_id=0 is the foreground class and class_id=1 is "__background__". Legacy keypoint checkpoints with args.num_keypoints_per_class[0] == 0 use a background-first layout: slot 0 maps to "__background__" and foreground slots map to class_names in order.

Note

A CPU tensor image is pinned before its transfer to a CUDA-device model, and that transfer is non-blocking; passing a tensor already on the model's accelerator skips this image transfer entirely. But with the default include_source_image=True, capturing source_image from that same tensor still does its own separate, blocking .cpu() call earlier in the loop — so an already-CUDA tensor input alone does not make the call fully round-trip-free. Pass include_source_image=False to avoid that copy as well. The exception is MPS with antialias=True on torch < 2.7, which has no MPS kernel for antialiased resize: there the whole preprocessing path, including an MPS-resident tensor input, runs on the CPU and the resized batch is transferred to the device once.

Tensor and non-uint8 NumPy range checks and every input's shape check are evaluated before inference. PIL images, natively decoded local files and uint8 NumPy images skip a redundant range scan because their byte-to-float conversion guarantees values in [0, 1]. Any resulting ValueError is raised only after all inputs have been inspected, so valid-shaped images later in a multi-image call still have their conversion and transfer queued before an earlier validation failure raises.

Raises:

Type Description
TypeError

If a source option of :meth:RFDETRInference.predict is passed, or if threshold is neither a number nor a sequence of numbers, such as a dict keyed by class name.

ValueError

If shape cannot be unpacked as a two-element sequence, if either dimension does not support the __index__ protocol (e.g. float) or is a bool, if either dimension is zero or negative, if either dimension is not divisible by patch_size * num_windows, if patch_size is not a positive integer, or if threshold is a sequence that is not one-dimensional, whose length differs from len(class_names), or that holds a NaN.

Examples:

Skipped because this example requires pretrained weights and a local image file.

>>> from rfdetr import RFDETRSmall
>>> model = RFDETRSmall()
>>> detections = model.predict("image.jpg", threshold=0.5)
Source code in src/rfdetr/detr.py
def predict(
    self,
    images: _ImageInput | Iterable[_ImageInput],
    threshold: _ThresholdInput = 0.5,
    shape: tuple[int, int] | None = None,
    patch_size: int | None = None,
    include_source_image: bool = True,
    *,
    antialias: bool = False,
    **kwargs: Any,
) -> _PredictionOutput:
    """Performs model inference on the input images.

    This method accepts a single image or an iterable of images (string or path-like file path, image URL string,
    PIL Image, NumPy array, or torch.Tensor) and runs them as one batch. The images should be in RGB channel order.
    A torch.Tensor must already be normalized to values in the [0, 1] range and have the shape (C, H, W); a
    (B, C, H, W) tensor is split into B images of one batch. Videos, directories, globs, manifests, webcams and
    streams are not images: predict them through ``RFDETRInference(model).predict(...)``.

    Args:
        images:
            A single image or an iterable of images to process. Images can be provided as string or path-like
            file paths, URL strings, PIL Images, NumPy arrays, or torch.Tensors. Iterables other than lists and
            tuples are read into a list first. PIL and file inputs convert to grayscale for one-channel models and
            RGB for three-channel models.
            HTTP(S) URL strings are downloaded with ``requests`` using a 30-second connect and read timeout,
            which does not bound the total download time. Redirects are followed, and the whole response body is
            read into memory with no size limit or host restriction. Pass only URLs you trust; for untrusted
            URLs, download and validate the image yourself and pass the decoded image instead.
        threshold:
            The confidence score a detection must exceed to be kept. A single float applies to every class. A
            :class:`~rfdetr.ClassThresholds`, such as ``ClassThresholds({"person": 0.35}, default=0.5)``, applies
            each class's own threshold to its detections, so that a class the model finds less reliably can be
            held to a different bar. It leaves out every detection that belongs to no entry of ``class_names``:
            the no-object slot of a fine-tuned model, the background slot of a legacy keypoint checkpoint, and the
            unused category IDs of a pretrained COCO model. A threshold of ``float("inf")`` drops every detection of
            that class.
        shape:
            Resize to ``(height, width)``; defaults to the model resolution. Both dimensions must be positive
            integers divisible by ``patch_size * num_windows``. Use the model's export resolution when applicable.
        patch_size:
            Backbone patch size for shape validation. Defaults to ``model_config.patch_size``.
        include_source_image:
            Whether to attach the original image to the returned prediction. Detection and segmentation outputs use
            ``detections.metadata["source_image"]``. Keypoint outputs use per-object
            ``key_points.data["source_image"]`` because Supervision ``KeyPoints`` currently has no collection-level
            metadata field. Defaults to ``True``. Set to ``False`` to reduce memory use when source images are not
            needed.
        antialias:
            Whether to use antialiasing during the inference resize. Match it to the resize the checkpoint was
            trained with. Keep this ``False`` for checkpoints trained with the antialias-free Albumentations
            pipeline, which is what the default CPU augmentation backend uses when ``rfdetr[augment]`` is
            installed. Set it to ``True`` for checkpoints trained with torchvision resizing: the Kornia/GPU
            augmentation backend, the CPU backend without ``rfdetr[augment]``, older RF-DETR releases, or
            antialiased platform preprocessing. Whether ``rfdetr[augment]`` was installed at training time
            therefore decides which setting a default-trained checkpoint expects.
            ``RFDETRInference`` uses the artifact metadata's ``resize_antialias`` by default for exported
            models; callers can override it with an explicit ``antialias`` argument. The lower-level
            ``rfdetr.export`` runtime preprocessing helpers always resize with ``antialias=False``.
            Defaults to ``False`` to preserve the existing inference behavior.
        **kwargs:
            Unknown keyword arguments are ignored with a FutureWarning. Passing them is deprecated since
            v1.12.0 and will raise TypeError in v2.0.0. The source options of
            :meth:`RFDETRInference.predict` (``stream``, ``batch``, ``vid_stride``, ``stream_buffer``) raise
            ``TypeError`` now.

    Returns:
        One result for a single image, or a list with one result per image for an iterable or a BCHW tensor.
        Detection and segmentation return :class:`~supervision.Detections`; keypoints return
        :class:`~supervision.KeyPoints` with ``xy`` coordinates.
        ``key_points.detection_confidence`` is the per-object score used by ``threshold``. For keypoint models this
        is the postprocessed detection score and, by default, includes normalized keypoint uncertainty fusion
        controlled by ``model_config.postprocess_trace_alpha``. ``key_points.keypoint_confidence`` is separate: it
        is a ``(num_detections, num_keypoints)`` array of per-keypoint findability scores decoded from the keypoint
        head, not a repeated copy of the detection score. When RF-DETR emits keypoint precision parameters,
        ``key_points.data["covariance"]`` stores per-keypoint pixel-space covariance matrices with shape
        ``(num_detections, num_keypoints, 2, 2)``. ``key_points.data["xyxy"]`` stores the corresponding detection
        boxes as a ``(num_detections, 4)`` array in the same row order as ``key_points.xy`` because Supervision
        ``KeyPoints`` does not have a native bounding-box field. The ``data`` dict also contains ``class_name`` and
        ``source_shape`` as per-object arrays. When ``include_source_image=True`` for keypoint models,
        ``source_image`` is stored as per-object data until Supervision exposes collection-level metadata for
        ``KeyPoints``.

    Note:
        For ``Detections`` outputs, ``source_image`` moved from ``detections.data`` to ``detections.metadata``.
        Update detection callers reading ``detections.data["source_image"]`` to use
        ``detections.metadata["source_image"]``.

    Note:
        ``class_name`` mapping uses one of three modes depending on the checkpoint. For pretrained COCO checkpoints
        (detected when ``model.args.num_classes > len(class_names)`` and ``class_names`` matches
        ``COCO_CLASS_NAMES``), raw COCO category IDs (1–90, sparse) are looked up by category ID rather than by
        position — so ``class_id=18`` yields ``"dog"``, not ``class_names[18]``. For fine-tuned detection and
        segmentation models and active-first keypoint models, ``class_id`` is a 0-based index into
        ``class_names``. In the one-class preview keypoint setup, that means ``class_id=0`` is the foreground
        class and ``class_id=1`` is ``"__background__"``.
        Legacy keypoint checkpoints with ``args.num_keypoints_per_class[0] == 0`` use a background-first layout:
        slot 0 maps to ``"__background__"`` and foreground slots map to ``class_names`` in order.

    Note:
        A CPU tensor image is pinned before its transfer to a CUDA-device model, and that transfer is
        non-blocking; passing a tensor already on the model's accelerator skips this image transfer entirely.
        But with the default ``include_source_image=True``, capturing ``source_image`` from that same tensor
        still does its own separate, blocking ``.cpu()`` call earlier in the loop — so an already-CUDA tensor
        input alone does not make the call fully round-trip-free. Pass ``include_source_image=False`` to avoid
        that copy as well. The exception is MPS with ``antialias=True`` on torch < 2.7, which has no MPS kernel
        for antialiased resize: there the whole preprocessing path, including an MPS-resident tensor input, runs
        on the CPU and the resized batch is transferred to the device once.

        Tensor and non-uint8 NumPy range checks and every input's shape check are evaluated before inference.
        PIL images, natively decoded local files and uint8 NumPy images skip a redundant range scan because
        their byte-to-float conversion guarantees values in ``[0, 1]``. Any resulting ``ValueError`` is raised
        only after all inputs have been inspected, so valid-shaped images later in a multi-image call still
        have their conversion and transfer queued before an earlier validation failure raises.

    Raises:
        TypeError: If a source option of :meth:`RFDETRInference.predict` is passed, or if ``threshold`` is neither a
            number nor a sequence of numbers, such as a ``dict`` keyed by class name.
        ValueError: If ``shape`` cannot be unpacked as a two-element sequence,
            if either dimension does not support the ``__index__`` protocol (e.g. ``float``) or is a ``bool``, if
            either dimension is zero or negative, if either dimension is not divisible by ``patch_size *
            num_windows``, if ``patch_size`` is not a positive integer, or if ``threshold`` is a sequence that is
            not one-dimensional, whose length differs from ``len(class_names)``, or that holds a NaN.

    Examples:
        Skipped because this example requires pretrained weights and a local image file.

        >>> from rfdetr import RFDETRSmall  # doctest: +SKIP
        >>> model = RFDETRSmall()  # doctest: +SKIP
        >>> detections = model.predict("image.jpg", threshold=0.5)  # doctest: +SKIP
    """
    source_options = sorted(_SOURCE_PREDICT_OPTIONS.intersection(kwargs))
    if source_options:
        raise TypeError(
            f"RFDETR.predict does not accept {', '.join(source_options)}: it predicts images only. "
            "Use RFDETRInference(model).predict(...) for videos, directories, globs, manifests, webcams, "
            "streams and lazy results."
        )
    if kwargs:
        deprecation_warning(
            "Passing unknown keyword arguments to RFDETR.predict is deprecated since v1.12.0: "
            f"{', '.join(sorted(kwargs))}. They are ignored and will raise TypeError in v2.0.0.",
            stacklevel=2,
        )
    if isinstance(images, torch.Tensor) and images.ndim == 4:
        images = list(images.unbind(0))
    return _predict_batch(
        self._prediction_context, images, threshold, shape, patch_size, include_source_image, antialias
    )

remove_optimized_model()

Clear legacy optimization state on this RFDETR instance.

Predictors returned by :meth:inference own independent runtimes; this method does not change or release them. Release the predictor itself when it is no longer needed. A legacy instance optimized destructively cannot recover its original module, so cleanup warns and leaves that state intact.

Source code in src/rfdetr/detr.py
def remove_optimized_model(self) -> None:
    """Clear legacy optimization state on this RFDETR instance.

    Predictors returned by :meth:`inference` own independent runtimes; this method does not change or release them.
    Release the predictor itself when it is no longer needed. A legacy instance optimized destructively cannot
    recover its original module, so cleanup warns and leaves that state intact.
    """
    if getattr(self, "_optimized_inplace", False):
        warnings.warn(
            "remove_optimized_model() has no effect after inplace optimization — the original model "
            "cannot be restored because export() and dtype casting mutate it. "
            "Create or reload a new RFDETR instance instead.",
            UserWarning,
            stacklevel=2,
        )
        return
    self.model.inference_model = None
    self._is_optimized_for_inference = False
    self._optimized_has_been_compiled = False
    self._optimized_batch_size = None
    self._optimized_resolution = None
    self._optimized_dtype = None
    self._optimized_inplace = False

train(**kwargs)

Train an RF-DETR model via the PyTorch Lightning stack.

All keyword arguments are forwarded to :meth:get_train_config to build a :class:~rfdetr.config.TrainConfig. Several kwargs are absorbed and handled specially so that existing call-sites do not break:

  • resolution — updates the model's input resolution by mutating :attr:model_config.resolution in place before the train config is built. This change persists on :attr:model_config after :meth:train returns. The value must be a positive integer divisible by patch_size * num_windows for the model variant; a :class:ValueError is raised otherwise. :attr:model_config.positional_encoding_size is also updated when the config derives it formulaically (PE == resolution // patch_size); configs with a pretrained-specific PE value (e.g. RFDETRBase uses DINOv2's PE=37 at 560 px) are left unchanged to preserve checkpoint compatibility.
  • device — normalized via :class:torch.device and mapped to PyTorch Lightning trainer arguments. "cpu" becomes accelerator="cpu"; "cuda" and "cuda:N" become accelerator="gpu" and optionally devices=[N]; "mps" becomes accelerator="mps"; "xla" becomes accelerator="tpu" (PTL's canonical name for the XLA backend). Other valid torch device types fall back to PTL auto-detection and emit a :class:UserWarning.
  • notes — optional user-defined metadata (string, dict, list, or any JSON-serialisable value) stored under the "notes" key in every .pth checkpoint produced during training. The value is also available inside args["notes"] for full provenance. Pass the same value to :meth:export to embed it in the exported artifact's metadata as well.

With resume pointing at a checkpoint that carries optimizer state, the optimizer's moments and the schedule's progress carry over, and the learning-rate and weight-decay settings (lr, lr_encoder, lr_vit_layer_decay, lr_component_decay, weight_decay) follow the order defaults, then the checkpoint, then the keyword arguments: one passed here wins, and one left out takes the value recorded in the checkpoint (#1613).

After training completes the underlying nn.Module is synced back onto self.model.model so that :meth:predict and :meth:export continue to work without reloading the checkpoint.

Raises:

Type Description
ImportError

If training dependencies are not installed. Install with pip install "rfdetr[train,loggers]".

ValueError

If resolution is not a positive integer or is not divisible by patch_size * num_windows for the model variant.

Source code in src/rfdetr/detr.py
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
def train(self, **kwargs: Any) -> None:
    """Train an RF-DETR model via the PyTorch Lightning stack.

    All keyword arguments are forwarded to :meth:`get_train_config` to build a :class:`~rfdetr.config.TrainConfig`.
    Several kwargs are absorbed and handled specially so that existing call-sites do not break:

    * ``resolution`` — updates the model's input resolution by mutating
      :attr:`model_config.resolution` in place before the train config is built. This change persists on
      :attr:`model_config` after :meth:`train` returns. The value must be a positive integer divisible by
      ``patch_size * num_windows`` for the model variant; a :class:`ValueError` is raised otherwise.
      :attr:`model_config.positional_encoding_size` is also updated when the config derives it formulaically (``PE
      == resolution // patch_size``); configs with a pretrained-specific PE value (e.g. ``RFDETRBase`` uses DINOv2's
      PE=37 at 560 px) are left unchanged to preserve checkpoint compatibility.
    * ``device`` — normalized via :class:`torch.device` and mapped to PyTorch
      Lightning trainer arguments. ``"cpu"`` becomes ``accelerator="cpu"``; ``"cuda"`` and ``"cuda:N"`` become
      ``accelerator="gpu"`` and optionally ``devices=[N]``; ``"mps"`` becomes ``accelerator="mps"``; ``"xla"``
      becomes ``accelerator="tpu"`` (PTL's canonical name for the XLA backend). Other valid torch device types
      fall back to PTL auto-detection and emit a :class:`UserWarning`.
    * ``notes`` — optional user-defined metadata (string, dict, list, or
      any JSON-serialisable value) stored under the ``"notes"`` key in every ``.pth`` checkpoint produced during
      training.  The value is also available inside ``args["notes"]`` for full provenance.  Pass the same value to
      :meth:`export` to embed it in the exported artifact's metadata as well.

    With ``resume`` pointing at a checkpoint that carries optimizer state, the optimizer's moments and the
    schedule's progress carry over, and the learning-rate and weight-decay settings (``lr``, ``lr_encoder``,
    ``lr_vit_layer_decay``, ``lr_component_decay``, ``weight_decay``) follow the order defaults, then the
    checkpoint, then the keyword arguments: one passed here wins, and one left out takes the value recorded in the
    checkpoint (#1613).

    After training completes the underlying ``nn.Module`` is synced back onto ``self.model.model`` so that
    :meth:`predict` and :meth:`export` continue to work without reloading the checkpoint.

    Raises:
        ImportError: If training dependencies are not installed. Install with
            ``pip install "rfdetr[train,loggers]"``.
        ValueError: If ``resolution`` is not a positive integer or is not
            divisible by ``patch_size * num_windows`` for the model variant.
    """
    # The training stack lives in the `rfdetr[train]` extras group — a missing
    # `pytorch_lightning` (or any other training-extras package) causes the import to fail,
    # and the remediation is `pip install "rfdetr[train,loggers]"`.
    try:
        from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
    except ModuleNotFoundError as exc:
        # Preserve internal import errors so packaging/regression issues in
        # rfdetr.* are not misreported as missing optional extras.
        if exc.name and exc.name.startswith("rfdetr."):
            raise
        raise ImportError(
            "RF-DETR training dependencies are missing. "
            'Install them with `pip install "rfdetr[train,loggers]"` and try again.',
        ) from exc

    if getattr(self, "_has_been_trained", False):
        warnings.warn(
            "Calling train() on a model that has already been trained or loaded from a checkpoint. "
            "The new training run will start from the original pretrained weights (pretrain_weights), "
            "NOT from the in-memory trained state. To continue training, pass resume=<checkpoint_path>.",
            UserWarning,
            stacklevel=2,
        )

    # Absorb the special train/evaluate kwargs (device, resolution, deprecated knobs),
    # build the TrainConfig, and resolve any auto batch size.  Shared with evaluate()
    # so both accept exactly the same keyword arguments.
    config, _accelerator, _devices = _prepare_run_config(self, **kwargs)
    # Before anything records the config: training_config.json and every checkpoint describe the resolved one.
    if config.resume:
        config = _resolve_resumed_optimizer_settings(config)

    # Auto-detect num_classes from the training dataset and align model_config.
    # This must run before RFDETRModelModule is constructed so that weight loading
    # inside the module uses the correct (dataset-derived) class count.
    dataset_dir = getattr(config, "dataset_dir", None)
    if dataset_dir:
        self._align_keypoint_schema_from_dataset(config)
        self._align_num_classes_from_dataset(dataset_dir)

    # Record the run's configuration before anything expensive starts, so a run interrupted while the model is
    # built, while dataset grids render, or during training itself still leaves one behind (#1493). Rewritten
    # after trainer.fit() with the dataset's class names. Both configs are final here: the alignment above is
    # the last thing that mutates them, and neither module construction nor build_trainer touches them.
    #
    # The dataset's class names are not known until the datamodule builds a dataset inside fit(), so read the
    # label space straight off disk instead: the COCO/YOLO readers first, then the train shard index of a
    # packed dataset_file="webdataset" directory. Best-effort on the same exception tuple as the num_classes
    # alignment above: a layout none of these understand records class_names: null until the post-fit write
    # fills it in, rather than blocking training.
    #
    # Guarded on the launcher's environment rather than is_main_process(), which reports rank 0 in every
    # process until trainer.fit() initializes torch.distributed; several ranks would otherwise write this one
    # path at once, where a torn write can truncate a previous run's good copy. Same guard as the dataset-grid
    # block below.
    # Reject a class_names that setup("fit") would reject inside trainer.fit(), before anything records or acts on
    # it: the write below would save the rejected list, and PTL's _call_setup_hook starts every configured
    # logger's experiment (wandb.init(), an MLflow run) before it calls datamodule setup, leaving an orphan
    # crashed run behind. Runs on every rank, so none is left hanging in fit()'s first collective; keypoint mode
    # is skipped for the reason the read below is -- the slot layout needs a built dataset.
    if (
        getattr(config, "class_names", None) is not None
        and dataset_dir
        and not self.model_config.use_grouppose_keypoints
    ):
        from rfdetr.training.module_data import _check_class_names_match_dataset

        if not hasattr(self, "_coco_categories_cache"):
            self._coco_categories_cache = {}
        try:
            dataset_class_names, labels_remapped = RFDETR._dataset_label_space_on_disk(
                dataset_dir,
                config.dataset_file,
                coco_categories=RFDETR._memoized_coco_categories(self._coco_categories_cache, dataset_dir),
            )
        except (FileNotFoundError, ValueError, KeyError, OSError) as exc:
            # Best-effort, like the read below: an unreadable layout fails with a better message inside fit().
            logger.debug("Could not read class names from dataset '%s': %s", dataset_dir, exc)
            dataset_class_names, labels_remapped = None, False
        # Outside the except above: the ValueError this raises is the verdict, not a failed read.
        _check_class_names_match_dataset(config.class_names, dataset_class_names, labels_remapped=labels_remapped)

    if _is_launcher_main_process():
        pre_fit_class_names = getattr(config, "class_names", None)
        # Keypoint mode stays null: the readers below return the detection basis (e.g. ['person']), but the
        # slot layout — a background-first schema such as [0, 17] pads a leading '' — is only known post-fit.
        if pre_fit_class_names is None and dataset_dir and not self.model_config.use_grouppose_keypoints:
            if not hasattr(self, "_coco_categories_cache"):
                self._coco_categories_cache = {}
            try:
                # Reuses the parse the num_classes alignment above just did, when the layout is COCO.
                pre_fit_class_names = RFDETR._load_classes(
                    dataset_dir,
                    coco_categories=RFDETR._memoized_coco_categories(self._coco_categories_cache, dataset_dir),
                )
            except (FileNotFoundError, ValueError, KeyError, OSError) as exc:
                logger.debug("Could not read class names from dataset '%s': %s", dataset_dir, exc)
                pre_fit_class_names = None
            if pre_fit_class_names is None and (Path(dataset_dir) / index_name("train")).exists():
                # A packed directory has no raw annotation file for _load_classes; its train shard index
                # carries the categories under the same "remap"/"raw" convention the num_classes alignment
                # above read. The tuple is wider than _detect_num_classes_for_training's own
                # WebDatasetSplitUnavailableError guard on purpose: a corrupt index raises ValueError/KeyError
                # from ShardIndex.from_json, and that must not block training either.
                try:
                    pre_fit_class_names = read_shard_index(dataset_dir, "train").class_names()
                except (FileNotFoundError, ValueError, KeyError, OSError) as exc:
                    logger.debug("Could not read class names from the shard index in '%s': %s", dataset_dir, exc)
        _save_training_config(config, self.model_config, pre_fit_class_names)
    else:
        # A non-zero launcher rank is the ordinary DDP case, but a non-zero task of a multi-task `srun` step
        # (SLURM_PROCID≠0) — e.g. a per-task sweep with devices=1, where every task is its own single-process
        # run — lands here too, writing neither this file nor the dataset grids. Say so at INFO, naming the
        # variable, rather than leaving an absent file with no explanation anywhere; WARNING would fire on
        # N-1 ranks of every ordinary DDP run.
        logger.info(
            "Not the launcher's main process (RANK=%s, SLURM_PROCID=%s); skipping the start-of-run "
            "training_config.json write.",
            os.environ.get("RANK"),
            os.environ.get("SLURM_PROCID"),
        )

    module = RFDETRModelModule(self.model_config, config)
    datamodule = RFDETRDataModule(self.model_config, config)

    # Guard on the launcher's environment rather than is_main_process() because torch.distributed is not yet
    # initialized here (it is set up inside trainer.fit()).  This used to read LOCAL_RANK alone, which let one
    # process per node through on multi-node runs and every process through under srun, which sets neither
    # LOCAL_RANK nor NODE_RANK.
    if config.save_dataset_grids and _is_launcher_main_process():
        # Outside the try: building the datasets is fatal to trainer.fit() below whatever happens here, so
        # catching its failure only relabels the real cause as a grid-save warning and then reports it twice.
        # The import stays inside, where a missing visualization dependency keeps costing only the grids.
        datamodule.setup("fit")
        try:
            from rfdetr.datasets.save_grids import DatasetGridSaver

            grids_output_dir = Path(config.output_dir) / "dataset_grids"
            DatasetGridSaver(datamodule.train_dataloader(), grids_output_dir, dataset_type="train").save_grid()
            DatasetGridSaver(datamodule.val_dataloader(), grids_output_dir, dataset_type="val").save_grid()
        except Exception:
            logger.warning(
                "Failed to save dataset grids; training will continue without them.",
                exc_info=True,
            )
    elif config.save_dataset_grids:
        logger.info(
            "Not the launcher's main process (RANK=%s, SLURM_PROCID=%s); skipping the dataset-grid render.",
            os.environ.get("RANK"),
            os.environ.get("SLURM_PROCID"),
        )

    if config.resume:
        # BestModelCallback's four lightweight checkpoint files (unlike the trainer's own
        # `last.ckpt` / `checkpoint_<epoch>.ckpt`, which retain full PTL state) intentionally
        # omit optimizer/LR-scheduler state to stay small — see
        # BestModelCallback._build_checkpoint_payload. They retain model weights, epoch
        # metadata, and per-callback state only when the resumed callback configuration
        # matches the saved keys. Checkpoints written before callback-state persistence have
        # no such state. Best-score tracking additionally requires exactly the original
        # output_dir because of PTL's ModelCheckpoint.load_state_dict() dirpath gate.
        # Flag the optimizer/scheduler gap explicitly instead of letting it pass silently.
        _light_checkpoint_names = frozenset(
            {"checkpoint_best_regular.pth", "checkpoint_best_ema.pth", "checkpoint_best_total.pth", "last_ema.pth"}
        )
        if Path(config.resume).name in _light_checkpoint_names:
            from rfdetr.utilities.io import _safe_torch_load

            # checkpoint.get("callbacks") is None-checked by PTL's own
            # _call_callbacks_load_state_dict(), which no-ops (skipping every callback's
            # restoration) when the key is absent or empty. Checkpoints written before
            # BestModelCallback._build_checkpoint_payload started persisting per-callback state
            # — or from a run where every registered callback happened to have empty state —
            # are exactly this case, so peek at the file rather than let the warning below
            # overclaim a restoration that silently does not happen.
            _resume_ckpt = _safe_torch_load(config.resume, trust=True)
            _has_callback_state = bool(_resume_ckpt.get("callbacks"))
            del _resume_ckpt
            _resume_dir = Path(config.resume).resolve().parent
            _configured_output_dir = Path(config.output_dir).resolve()
            _best_score_restores = _resume_dir == _configured_output_dir

            if _has_callback_state:
                logger.warning(
                    "resume=%r points at one of BestModelCallback's lightweight checkpoints, "
                    "which intentionally omit optimizer/LR-scheduler state to stay small. "
                    "Model weights and epoch count will resume. Callback state can restore only "
                    "for matching configured callbacks; the optimizer and LR scheduler restart cold. "
                    "To resume with optimizer/scheduler state too, pass the trainer's full "
                    "checkpoint instead (e.g. %s/last.ckpt or %s/checkpoint_<epoch>.ckpt).",
                    config.resume,
                    config.output_dir,
                    config.output_dir,
                )
            else:
                logger.warning(
                    "resume=%r points at one of BestModelCallback's lightweight checkpoints, "
                    "which intentionally omit optimizer/LR-scheduler state to stay small. "
                    "Model weights and epoch count will resume, but this particular file has no "
                    "saved callback state (it predates callback-state persistence, or every "
                    "registered callback had nothing to save), so best-score tracking, EMA, and "
                    "early-stopping state all restart cold this run too — not just the "
                    "optimizer and LR scheduler. To resume with full state, pass the trainer's "
                    "full checkpoint instead (e.g. %s/last.ckpt or %s/checkpoint_<epoch>.ckpt).",
                    config.resume,
                    config.output_dir,
                    config.output_dir,
                )

            # BestModelCallback always writes these four files directly under its own
            # `dirpath` (== output_dir at save time; see BestModelCallback.__init__ and
            # _build_checkpoint_payload). PTL's ModelCheckpoint.load_state_dict() only
            # restores best_model_score/best_k_models/kth_value/last_model_path when the
            # resumed dirpath matches the checkpoint's saved dirpath exactly (model_checkpoint.py,
            # installed pytorch-lightning) — with a different output_dir this run only recovers
            # best_model_path, so the first metric logged after resume looks like an automatic
            # improvement over an empty best_model_score. Only worth flagging when there was
            # callback state to lose in the first place.
            if _has_callback_state and not _best_score_restores:
                logger.warning(
                    "resume=%r was written under %s but output_dir=%r points elsewhere. "
                    "PyTorch Lightning only restores best_model_score/best_k_models when "
                    "output_dir matches the checkpoint's original directory exactly — with "
                    "this output_dir, best-score tracking (BestModelCallback's high-water "
                    "mark) restarts fresh in the new directory instead of resuming. Set "
                    "output_dir=%r to keep it.",
                    config.resume,
                    _resume_dir,
                    config.output_dir,
                    str(_resume_dir),
                )

    trainer_kwargs: dict[str, Any] = {"accelerator": _accelerator}
    if _devices is not None:
        trainer_kwargs["devices"] = _devices
    trainer = build_trainer(config, self.model_config, **trainer_kwargs)
    trainer.fit(module, datamodule, ckpt_path=config.resume or None)

    # Sync the trained weights back so predict() / export() see the updated model.
    self.model.model = module.model
    # Rebuild model.args from the real training config (#1199): it was previously left as the
    # construction-time snapshot built from a dummy TrainConfig, so overrides like lr/lr_encoder
    # never appeared in model.model.__dict__['args'] even though the optimizer used them correctly.
    self.model.args = _namespace_from_configs(self.model_config, config)
    # Mark this instance as trained so a subsequent train() call warns that it will
    # restart from pretrain_weights rather than continue from the in-memory state.
    self._has_been_trained = True
    # Invalidate any compiled inference snapshot: it was built from the pre-training
    # weights and must not survive the model reassignment above.
    self.remove_optimized_model()
    # Sync class names: prefer explicit config.class_names, otherwise fall back to dataset (#509).
    config_class_names = getattr(config, "class_names", None)
    if config_class_names is not None:
        self.model.class_names = config_class_names
    else:
        dataset_class_names = getattr(datamodule, "class_names", None)
        if dataset_class_names is not None:
            self.model.class_names = dataset_class_names

    # Rewrite the configuration saved at the start of the run, now that the dataset's class names are known.
    # Guard to main process only to avoid races in distributed/multi-GPU training; unlike the start-of-run
    # write before trainer.fit(), torch.distributed is initialized here, so the global rank is authoritative.
    if is_main_process():
        _save_training_config(config, self.model_config, self.model.class_names)

write_export_metadata(artifact, *, format='onnx', shape=None, patch_size=None, batch_size=1, dynamic_batch=False, max_batch_size=None, backbone_only=False, backend=None, soc=None, **export_settings)

Write inference metadata for an existing artifact from its live source model.

Load the checkpoint the artifact was exported from and supply the same export settings. This method reads model configuration and prediction semantics without running the model or converting the artifact. Each format's exporter supplies its artifact-specific input and output mapping before metadata is written.

Parameters:

Name Type Description Default

artifact

str | PathLike[str]

Existing exported file or model bundle.

required

format

str

Export format, with the same names and aliases accepted by :meth:export.

'onnx'

shape

tuple[int, int] | None

Exported (height, width); defaults to the source model's square resolution.

None

patch_size

int | None

Backbone patch size; must match the source model when supplied.

None

batch_size

int

Positive batch size used to export the artifact.

1

dynamic_batch

bool

Whether the artifact was exported with a dynamic batch dimension.

False

max_batch_size

int | None

TensorRT dynamic-batch profile upper bound; ignored for other exports.

None

backbone_only

bool

Whether the artifact contains only backbone features.

False

backend

str | None

ExecuTorch backend used to export the artifact.

None

soc

str | None

Target Qualcomm SoC for the ExecuTorch QNN backend.

None

**export_settings

Any

Format settings forwarded to the exporter's configuration builder, such as coreai_precision. Use the values supplied to :meth:export.

{}

Returns:

Type Description
Path | None

Adjacent digest-bound metadata path, or None when metadata is embedded in an ONNX artifact.

Raises:

Type Description
FileNotFoundError

If the artifact does not exist.

ValueError

If the format, batch bounds, shape, patch size, backend, SoC, or format settings are invalid.

NotImplementedError

If the requested settings are unsupported by the format.

ImportError

If embedding metadata or reading the artifact's interface needs a missing optional package.

OSError

If the artifact cannot be read or metadata cannot be written.

Examples:

A real source checkpoint and its exported artifact are required.

>>> from rfdetr import RFDETRSmall
>>> model = RFDETRSmall(pretrain_weights="checkpoint.pth")
>>> model.write_export_metadata("output/model.xml", format="openvino")
Source code in src/rfdetr/detr.py
def write_export_metadata(
    self,
    artifact: str | os.PathLike[str],
    *,
    format: str = "onnx",
    shape: tuple[int, int] | None = None,
    patch_size: int | None = None,
    batch_size: int = 1,
    dynamic_batch: bool = False,
    max_batch_size: int | None = None,
    backbone_only: bool = False,
    backend: str | None = None,
    soc: str | None = None,
    **export_settings: Any,
) -> Path | None:
    """Write inference metadata for an existing artifact from its live source model.

    Load the checkpoint the artifact was exported from and supply the same export settings. This method reads
    model configuration and prediction semantics without running the model or converting the artifact. Each
    format's exporter supplies its artifact-specific input and output mapping before metadata is written.

    Args:
        artifact: Existing exported file or model bundle.
        format: Export format, with the same names and aliases accepted by :meth:`export`.
        shape: Exported ``(height, width)``; defaults to the source model's square resolution.
        patch_size: Backbone patch size; must match the source model when supplied.
        batch_size: Positive batch size used to export the artifact.
        dynamic_batch: Whether the artifact was exported with a dynamic batch dimension.
        max_batch_size: TensorRT dynamic-batch profile upper bound; ignored for other exports.
        backbone_only: Whether the artifact contains only backbone features.
        backend: ExecuTorch backend used to export the artifact.
        soc: Target Qualcomm SoC for the ExecuTorch QNN backend.
        **export_settings: Format settings forwarded to the exporter's configuration builder, such as
            ``coreai_precision``. Use the values supplied to :meth:`export`.

    Returns:
        Adjacent digest-bound metadata path, or ``None`` when metadata is embedded in an ONNX artifact.

    Raises:
        FileNotFoundError: If the artifact does not exist.
        ValueError: If the format, batch bounds, shape, patch size, backend, SoC, or format settings are invalid.
        NotImplementedError: If the requested settings are unsupported by the format.
        ImportError: If embedding metadata or reading the artifact's interface needs a missing optional package.
        OSError: If the artifact cannot be read or metadata cannot be written.

    Examples:
        A real source checkpoint and its exported artifact are required.

        >>> from rfdetr import RFDETRSmall
        >>> model = RFDETRSmall(pretrain_weights="checkpoint.pth")  # doctest: +SKIP
        >>> model.write_export_metadata("output/model.xml", format="openvino")  # doctest: +SKIP
    """
    from rfdetr.export._backend import _resolve_export_backend
    from rfdetr.export._runtime.metadata import metadata_from_model, write_metadata
    from rfdetr.export.base import reject_unsupported_dynamic_batch
    from rfdetr.export.prepare import validate_export_shape
    from rfdetr.export.registry import resolve_exporter

    artifact = Path(artifact)
    if not artifact.exists():
        raise FileNotFoundError(artifact)
    format, batch_size, max_batch_size = RFDETR._resolve_export_batch(
        format, batch_size, dynamic_batch, max_batch_size
    )
    backend, soc = _resolve_export_backend(format, backend, soc)
    reject_unsupported_dynamic_batch(format, dynamic_batch=dynamic_batch)
    shape = validate_export_shape(
        shape,
        patch_size,
        self.model_config,
        self.model.resolution,
        resolve_patch_size=_resolve_patch_size,
        validate_shape_dims=_validate_shape_dims,
    )
    exporter_class = resolve_exporter(format)
    config = exporter_class.build_config(
        **export_settings,
        backbone_only=backbone_only,
        dynamic_batch=dynamic_batch,
        batch_size=batch_size,
        max_batch_size=max_batch_size,
        backend=backend,
        soc=soc,
    )
    exporter = exporter_class(config)
    metadata = metadata_from_model(
        self,
        format=format,
        shape=shape,
        batch_size=batch_size,
        dynamic_batch=dynamic_batch,
        backbone_only=backbone_only,
        backend=backend,
        max_batch_size=max_batch_size,
    )
    return write_metadata(artifact, exporter._metadata_for_artifact(metadata, artifact))