# ------------------------------------------------------------------------ # RF-DETR # Copyright (c) 2025 Roboflow. All Rights Reserved. # Licensed under the Apache License, Version 2.0 [see LICENSE for details] # ------------------------------------------------------------------------ from __future__ import annotations import json import math import os import types from typing import Any, cast import torch import torch.nn.functional as F # noqa: N812 from torch import Tensor, nn from transformers import AutoBackbone from rfdetr.models.backbone.dinov2_with_windowed_attn import ( WindowedDinov2WithRegistersBackbone, WindowedDinov2WithRegistersConfig, ) from rfdetr.utilities.logger import get_logger logger = get_logger() size_to_width = { "tiny": 192, "small": 384, "base": 768, "large": 1024, } size_to_config = { "small": "dinov2_small.json", "base": "dinov2_base.json", "large": "dinov2_large.json", } size_to_config_with_registers = { "small": "dinov2_with_registers_small.json", "base": "dinov2_with_registers_base.json", "large": "dinov2_with_registers_large.json", } def get_config(size: str, use_registers: bool) -> dict[str, Any]: config_dict = size_to_config_with_registers if use_registers else size_to_config current_dir = os.path.dirname(os.path.abspath(__file__)) configs_dir = os.path.join(current_dir, "dinov2_configs") config_path = os.path.join(configs_dir, config_dict[size]) with open(config_path) as f: dino_config: dict[str, Any] = json.load(f) return dino_config def compute_window_block_indexes( out_feature_indexes: list[int], window_block_indexes: list[int] | None = None ) -> list[int]: """Derive which 0-indexed encoder blocks use windowed self-attention (the rest use global). When `window_block_indexes` is provided explicitly, it is returned unchanged — this is the override path for callers that want to pin an exact routing (e.g. to match a published architecture spec) without disturbing the derived default used by existing configs/checkpoints. This override is only wired through `ModelDefaults` (`rfdetr/models/_defaults.py`) or direct `Backbone`/`DinoV2` construction — not the public pydantic `RFDETR*Config` classes, which set `ConfigDict(extra="forbid")` and raise on unknown fields. When omitted, the windowed set is derived as the complement of `out_feature_indexes` (treated as raw block indices) within `range(0, out_feature_indexes[-1] + 1)`. This mirrors the historical behavior relied upon by all released RF-DETR checkpoints; it does not renumber `out_feature_indexes` (1-indexed HF stage numbers) to 0-indexed block indices, so it is not guaranteed to match any particular paper-specified layer schedule. The derivation assumes `out_feature_indexes` is given in ascending order: it uses `out_feature_indexes[-1]` (not `max(out_feature_indexes)`) as the upper bound, so an unsorted list whose maximum is not last would under-cover the intended range. Args: out_feature_indexes: Encoder stage numbers whose feature maps are exported to the decoder. window_block_indexes: Explicit windowed-block override. `None` derives from `out_feature_indexes` using the legacy formula. Returns: 0-indexed encoder block positions that should run windowed attention. Examples: >>> compute_window_block_indexes([3, 6, 9, 12]) [0, 1, 2, 4, 5, 7, 8, 10, 11] >>> compute_window_block_indexes([3, 6, 9, 12], window_block_indexes=[0, 1, 3, 4, 6, 7, 9, 10]) [0, 1, 3, 4, 6, 7, 9, 10] """ if window_block_indexes is not None: return window_block_indexes window_block_indexes_set = set(range(out_feature_indexes[-1] + 1)) window_block_indexes_set.difference_update(out_feature_indexes) return sorted(window_block_indexes_set) class DinoV2(nn.Module): def __init__( self, shape: tuple[int, int] = (640, 640), out_feature_indexes: list[int] | None = None, size: str = "base", use_registers: bool = True, use_windowed_attn: bool = True, gradient_checkpointing: bool = False, load_dinov2_weights: bool = True, patch_size: int = 14, num_windows: int = 4, positional_encoding_size: int = 37, drop_path_rate: float = 0.0, window_block_indexes: list[int] | None = None, ) -> None: super().__init__() if out_feature_indexes is None: out_feature_indexes = [2, 4, 5, 9] name = f"facebook/dinov2-with-registers-{size}" if use_registers else f"facebook/dinov2-{size}" self.shape = shape self.patch_size = patch_size self.num_windows = num_windows # Create the encoder if not use_windowed_attn: assert not gradient_checkpointing, "Gradient checkpointing is not supported for non-windowed attention" assert load_dinov2_weights, "Using non-windowed attention requires loading dinov2 weights from hub" if drop_path_rate > 0.0: logger.warning( "drop_path_rate > 0.0 is not supported for non-windowed DinoV2 backbones." " drop_path will be ignored." ) self.encoder = AutoBackbone.from_pretrained( # type: ignore[no-untyped-call] name, out_features=[f"stage{i}" for i in out_feature_indexes], return_dict=False, ) else: dino_config = get_config(size, use_registers) num_hidden_layers = dino_config["num_hidden_layers"] if not out_feature_indexes: raise ValueError("out_feature_indexes must be non-empty.") if window_block_indexes is not None and ( len(set(window_block_indexes)) != len(window_block_indexes) or any(not (0 <= idx < num_hidden_layers) for idx in window_block_indexes) ): raise ValueError( f"window_block_indexes entries must be unique and within " f"[0, {num_hidden_layers}); got {window_block_indexes}." ) window_block_indexes = compute_window_block_indexes(out_feature_indexes, window_block_indexes) dino_config["return_dict"] = False dino_config["out_features"] = [f"stage{i}" for i in out_feature_indexes] dino_config["drop_path_rate"] = drop_path_rate implied_resolution = positional_encoding_size * patch_size if implied_resolution != dino_config["image_size"]: logger.warning( "Using a different number of positional encodings than DINOv2, which means" " we're not loading DINOv2 backbone weights. This is not a problem if" " finetuning a pretrained RF-DETR model." ) dino_config["image_size"] = implied_resolution load_dinov2_weights = False if patch_size != 14: logger.warning( f"Using patch size {patch_size} instead of 14, which means we're not loading" " DINOv2 backbone weights. This is not a problem if finetuning a pretrained" " RF-DETR model." ) dino_config["patch_size"] = patch_size load_dinov2_weights = False if use_registers: windowed_dino_config = WindowedDinov2WithRegistersConfig( **dino_config, num_windows=num_windows, window_block_indexes=window_block_indexes, gradient_checkpointing=gradient_checkpointing, ) else: windowed_dino_config = WindowedDinov2WithRegistersConfig( **dino_config, num_windows=num_windows, window_block_indexes=window_block_indexes, num_register_tokens=0, gradient_checkpointing=gradient_checkpointing, ) self.encoder = ( WindowedDinov2WithRegistersBackbone.from_pretrained( name, config=windowed_dino_config, ) if load_dinov2_weights else WindowedDinov2WithRegistersBackbone(windowed_dino_config) ) self._out_feature_channels = [size_to_width[size]] * len(out_feature_indexes) self._export = False def export(self) -> None: if self._export: return self._export = True shape = self.shape def make_new_interpolated_pos_encoding( position_embeddings: Tensor, patch_size: int, height: int, width: int ) -> Tensor: num_positions = position_embeddings.shape[1] - 1 dim = position_embeddings.shape[-1] height = height // patch_size width = width // patch_size class_pos_embed = position_embeddings[:, 0] patch_pos_embed = position_embeddings[:, 1:] # Reshape and permute patch_pos_embed = patch_pos_embed.reshape( 1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim ) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) # Use bicubic interpolation, disabling antialias only on MPS devices patch_pos_embed = F.interpolate( patch_pos_embed, size=(height, width), mode="bicubic", align_corners=False, antialias=patch_pos_embed.device.type != "mps", ) # Reshape back patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).reshape(1, -1, dim) return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) # If the shape of self.encoder.embeddings.position_embeddings # matches the shape of your new tensor, use copy_: with torch.no_grad(): new_positions = make_new_interpolated_pos_encoding( self.encoder.embeddings.position_embeddings, self.encoder.config.patch_size, shape[0], shape[1], ) # Create a new Parameter with the new size old_interpolate_pos_encoding = self.encoder.embeddings.interpolate_pos_encoding def new_interpolate_pos_encoding(self_mod: Any, embeddings: Tensor, height: int, width: int) -> Tensor: num_patches = embeddings.shape[1] - 1 num_positions = self_mod.position_embeddings.shape[1] - 1 # The precomputed table is valid only for this exact static export grid. if num_patches == num_positions and (height, width) == shape: return cast(Tensor, self_mod.position_embeddings) return cast(Tensor, old_interpolate_pos_encoding(embeddings, height, width)) self.encoder.embeddings.position_embeddings = nn.Parameter(new_positions) self.encoder.embeddings.interpolate_pos_encoding = types.MethodType( new_interpolate_pos_encoding, self.encoder.embeddings ) def forward(self, x: Tensor) -> list[Tensor]: block_size = self.patch_size * self.num_windows assert x.shape[2] % block_size == 0 and x.shape[3] % block_size == 0, ( f"Backbone requires input shape to be divisible by {block_size}, but got {x.shape}" ) x = self.encoder(x) return list(x[0]) if __name__ == "__main__": model = DinoV2() model.export() x = torch.randn(1, 3, 640, 640) logger.info(model(x)) for j in model(x): logger.info(j.shape)