inference-models 0.18.3__py3-none-any.whl

inference_models/models/rfdetr/rfdetr_backbone_pytorch.py

@@ -0,0 +1,394 @@
+import json
+import math
+import os
+import types
+
+import torch
+import torch.nn.functional as F
+from peft import PeftModel
+from torch import nn
+from transformers import AutoBackbone
+
+from inference_models.logger import LOGGER
+from inference_models.models.rfdetr.dinov2_with_windowed_attn import (
+    WindowedDinov2WithRegistersBackbone,
+    WindowedDinov2WithRegistersConfig,
+)
+from inference_models.models.rfdetr.misc import NestedTensor
+from inference_models.models.rfdetr.projector import MultiScaleProjector
+
+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, use_registers):
+    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, "r") as f:
+        dino_config = json.load(f)
+    return dino_config
+
+
+class DinoV2(nn.Module):
+    def __init__(
+        self,
+        shape=(640, 640),
+        out_feature_indexes=[2, 4, 5, 9],
+        size="base",
+        use_registers=True,
+        use_windowed_attn=True,
+        gradient_checkpointing=False,
+        load_dinov2_weights=True,
+        patch_size=14,
+        num_windows=4,
+        positional_encoding_size=37,
+    ):
+        super().__init__()
+
+        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"
+            self.encoder = AutoBackbone.from_pretrained(
+                name,
+                out_features=[f"stage{i}" for i in out_feature_indexes],
+                return_dict=False,
+            )
+        else:
+            window_block_indexes = set(range(out_feature_indexes[-1] + 1))
+            window_block_indexes.difference_update(out_feature_indexes)
+            window_block_indexes = list(window_block_indexes)
+
+            dino_config = get_config(size, use_registers)
+
+            dino_config["return_dict"] = False
+            dino_config["out_features"] = [f"stage{i}" for i in out_feature_indexes]
+            implied_resolution = positional_encoding_size * patch_size
+
+            if implied_resolution != dino_config["image_size"]:
+                LOGGER.warning(
+                    f"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):
+        if self._export:
+            return
+        self._export = True
+        shape = self.shape
+
+        def make_new_interpolated_pos_encoding(
+            position_embeddings, patch_size, height, width
+        ):
+
+            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 bilinear interpolation without antialias
+            patch_pos_embed = F.interpolate(
+                patch_pos_embed,
+                size=(height, width),
+                mode="bicubic",
+                align_corners=False,
+                antialias=True,
+            )
+
+            # 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, embeddings, height, width):
+            num_patches = embeddings.shape[1] - 1
+            num_positions = self_mod.position_embeddings.shape[1] - 1
+            if num_patches == num_positions and height == width:
+                return self_mod.position_embeddings
+            return 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):
+        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])
+
+
+class BackboneBase(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def get_named_param_lr_pairs(self, args, prefix: str):
+        raise NotImplementedError
+
+
+class Backbone(BackboneBase):
+    """backbone."""
+
+    def __init__(
+        self,
+        name: str,
+        pretrained_encoder: str = None,
+        window_block_indexes: list = None,
+        drop_path=0.0,
+        out_channels=256,
+        out_feature_indexes: list = None,
+        projector_scale: list = None,
+        use_cls_token: bool = False,
+        freeze_encoder: bool = False,
+        layer_norm: bool = False,
+        target_shape: tuple[int, int] = (640, 640),
+        rms_norm: bool = False,
+        backbone_lora: bool = False,
+        gradient_checkpointing: bool = False,
+        load_dinov2_weights: bool = True,
+        patch_size: int = 14,
+        num_windows: int = 4,
+        positional_encoding_size: bool = False,
+    ):
+        super().__init__()
+        # an example name here would be "dinov2_base" or "dinov2_registers_windowed_base"
+        # if "registers" is in the name, then use_registers is set to True, otherwise it is set to False
+        # similarly, if "windowed" is in the name, then use_windowed_attn is set to True, otherwise it is set to False
+        # the last part of the name should be the size
+        # and the start should be dinov2
+        name_parts = name.split("_")
+        assert name_parts[0] == "dinov2"
+        size = name_parts[-1]
+        use_registers = False
+        if "registers" in name_parts:
+            use_registers = True
+            name_parts.remove("registers")
+        use_windowed_attn = False
+        if "windowed" in name_parts:
+            use_windowed_attn = True
+            name_parts.remove("windowed")
+        assert (
+            len(name_parts) == 2
+        ), "name should be dinov2, then either registers, windowed, both, or none, then the size"
+        self.encoder = DinoV2(
+            size=name_parts[-1],
+            out_feature_indexes=out_feature_indexes,
+            shape=target_shape,
+            use_registers=use_registers,
+            use_windowed_attn=use_windowed_attn,
+            gradient_checkpointing=gradient_checkpointing,
+            load_dinov2_weights=load_dinov2_weights,
+            patch_size=patch_size,
+            num_windows=num_windows,
+            positional_encoding_size=positional_encoding_size,
+        )
+        # build encoder + projector as backbone module
+        if freeze_encoder:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+
+        self.projector_scale = projector_scale
+        assert len(self.projector_scale) > 0
+        # x[0]
+        assert (
+            sorted(self.projector_scale) == self.projector_scale
+        ), "only support projector scale P3/P4/P5/P6 in ascending order."
+        level2scalefactor = dict(P3=2.0, P4=1.0, P5=0.5, P6=0.25)
+        scale_factors = [level2scalefactor[lvl] for lvl in self.projector_scale]
+
+        self.projector = MultiScaleProjector(
+            in_channels=self.encoder._out_feature_channels,
+            out_channels=out_channels,
+            scale_factors=scale_factors,
+            layer_norm=layer_norm,
+            rms_norm=rms_norm,
+        )
+
+        self._export = False
+
+    def export(self):
+        self._export = True
+        self._forward_origin = self.forward
+        self.forward = self.forward_export
+
+        if isinstance(self.encoder, PeftModel):
+            LOGGER.info("Merging and unloading LoRA weights")
+            self.encoder.merge_and_unload()
+
+    def forward(self, tensor_list: NestedTensor):
+        """ """
+        # (H, W, B, C)
+        feats = self.encoder(tensor_list.tensors)
+        feats = self.projector(feats)
+        # x: [(B, C, H, W)]
+        out = []
+        for feat in feats:
+            m = tensor_list.mask
+            assert m is not None
+            mask = F.interpolate(m[None].float(), size=feat.shape[-2:]).to(torch.bool)[
+                0
+            ]
+            out.append(NestedTensor(feat, mask))
+        return out
+
+    def forward_export(self, tensors: torch.Tensor):
+        feats = self.encoder(tensors)
+        feats = self.projector(feats)
+        out_feats = []
+        out_masks = []
+        for feat in feats:
+            # x: [(B, C, H, W)]
+            b, _, h, w = feat.shape
+            out_masks.append(
+                torch.zeros((b, h, w), dtype=torch.bool, device=feat.device)
+            )
+            out_feats.append(feat)
+        return out_feats, out_masks
+
+    def get_named_param_lr_pairs(self, args, prefix: str = "backbone.0"):
+        num_layers = args.out_feature_indexes[-1] + 1
+        backbone_key = "backbone.0.encoder"
+        named_param_lr_pairs = {}
+        for n, p in self.named_parameters():
+            n = prefix + "." + n
+            if backbone_key in n and p.requires_grad:
+                lr = (
+                    args.lr_encoder
+                    * get_dinov2_lr_decay_rate(
+                        n,
+                        lr_decay_rate=args.lr_vit_layer_decay,
+                        num_layers=num_layers,
+                    )
+                    * args.lr_component_decay**2
+                )
+                wd = args.weight_decay * get_dinov2_weight_decay_rate(n)
+                named_param_lr_pairs[n] = {
+                    "params": p,
+                    "lr": lr,
+                    "weight_decay": wd,
+                }
+        return named_param_lr_pairs
+
+
+def get_dinov2_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12):
+    """
+    Calculate lr decay rate for different ViT blocks.
+
+    Args:
+        name (string): parameter name.
+        lr_decay_rate (float): base lr decay rate.
+        num_layers (int): number of ViT blocks.
+    Returns:
+        lr decay rate for the given parameter.
+    """
+    layer_id = num_layers + 1
+    if name.startswith("backbone"):
+        if "embeddings" in name:
+            layer_id = 0
+        elif ".layer." in name and ".residual." not in name:
+            layer_id = int(name[name.find(".layer.") :].split(".")[2]) + 1
+    return lr_decay_rate ** (num_layers + 1 - layer_id)
+
+
+def get_dinov2_weight_decay_rate(name, weight_decay_rate=1.0):
+    if (
+        ("gamma" in name)
+        or ("pos_embed" in name)
+        or ("rel_pos" in name)
+        or ("bias" in name)
+        or ("norm" in name)
+        or ("embeddings" in name)
+    ):
+        weight_decay_rate = 0.0
+    return weight_decay_rate