birder 0.2.3__py3-none-any.whl → 0.3.0__py3-none-any.whl

birder/net/detection/fcos.py

@@ -455,13 +455,8 @@ class FCOS(DetectionBaseNet):
 
             # Non-maximum suppression
             if self.soft_nms is not None:
-                # Actually much faster on CPU
-                device = image_boxes.device
-                (soft_scores, keep) = self.soft_nms(
-                    image_boxes.cpu(), image_scores.cpu(), image_labels.cpu(), score_threshold=0.001
-                )
-                keep = keep.to(device)
-                image_scores[keep] = soft_scores.to(device)
+                (soft_scores, keep) = self.soft_nms(image_boxes, image_scores, image_labels, score_threshold=0.001)
+                image_scores[keep] = soft_scores
             else:
                 keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh)
 

birder/net/detection/retinanet.py

@@ -417,13 +417,8 @@ class RetinaNet(DetectionBaseNet):
 
             # Non-maximum suppression
             if self.soft_nms is not None:
-                # Actually much faster on CPU
-                device = image_boxes.device
-                (soft_scores, keep) = self.soft_nms(
-                    image_boxes.cpu(), image_scores.cpu(), image_labels.cpu(), score_threshold=0.001
-                )
-                keep = keep.to(device)
-                image_scores[keep] = soft_scores.to(device)
+                (soft_scores, keep) = self.soft_nms(image_boxes, image_scores, image_labels, score_threshold=0.001)
+                image_scores[keep] = soft_scores
             else:
                 keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh)
 

birder/net/detection/rt_detr_v1.py

@@ -1087,6 +1087,7 @@ class RT_DETR_v1(DetectionBaseNet):
 
         return (detections, losses)
 
+    @torch.no_grad()  # type: ignore[untyped-decorator]
     def reparameterize_model(self) -> None:
         if self.reparameterized is True:
             return

birder/net/efficientformer_v1.py

@@ -357,16 +357,22 @@ class EfficientFormer_v1(BaseNet):
         resolution = (int(new_size[0] / (2**5)), int(new_size[1] / (2**5)))
         for m in self.body.modules():
             if isinstance(m, Attention):
-                m.attention_biases = nn.Parameter(
-                    interpolate_attention_bias(m.attention_biases, old_resolution, resolution)
-                )
+                with torch.no_grad():
+                    m.attention_biases = nn.Parameter(
+                        interpolate_attention_bias(m.attention_biases, old_resolution, resolution)
+                    )
 
-                pos = torch.stack(
-                    torch.meshgrid(torch.arange(resolution[0]), torch.arange(resolution[1]), indexing="ij")
-                ).flatten(1)
-                rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
-                rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
-                m.attention_bias_idxs = nn.Buffer(rel_pos)
+                    device = m.attention_biases.device
+                    pos = torch.stack(
+                        torch.meshgrid(
+                            torch.arange(resolution[0], device=device),
+                            torch.arange(resolution[1], device=device),
+                            indexing="ij",
+                        )
+                    ).flatten(1)
+                    rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
+                    rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
+                    m.attention_bias_idxs = nn.Buffer(rel_pos)
 
 
 registry.register_model_config(

birder/net/efficientformer_v2.py

@@ -554,26 +554,30 @@ class EfficientFormer_v2(DetectorBackbone):
                         attn.N = attn.resolution[0] * attn.resolution[1]
                         attn.N2 = attn.resolution2[0] * attn.resolution2[1]
 
-                        # Interpolate attention_biases
-                        attn.attention_biases = nn.Parameter(
-                            interpolate_attention_bias(attn.attention_biases, old_base, new_base)
-                        )
-
-                        k_pos = torch.stack(
-                            torch.meshgrid(
-                                torch.arange(attn.resolution[0]), torch.arange(attn.resolution[1]), indexing="ij"
+                        with torch.no_grad():
+                            # Interpolate attention_biases
+                            attn.attention_biases = nn.Parameter(
+                                interpolate_attention_bias(attn.attention_biases, old_base, new_base)
                             )
-                        ).flatten(1)
-                        q_pos = torch.stack(
-                            torch.meshgrid(
-                                torch.arange(0, attn.resolution[0], step=2),
-                                torch.arange(0, attn.resolution[1], step=2),
-                                indexing="ij",
-                            )
-                        ).flatten(1)
-                        rel_pos = (q_pos[..., :, None] - k_pos[..., None, :]).abs()
-                        rel_pos = (rel_pos[0] * attn.resolution[1]) + rel_pos[1]
-                        attn.attention_bias_idxs = nn.Buffer(torch.LongTensor(rel_pos), persistent=False)
+
+                            device = attn.attention_biases.device
+                            k_pos = torch.stack(
+                                torch.meshgrid(
+                                    torch.arange(attn.resolution[0], device=device),
+                                    torch.arange(attn.resolution[1], device=device),
+                                    indexing="ij",
+                                )
+                            ).flatten(1)
+                            q_pos = torch.stack(
+                                torch.meshgrid(
+                                    torch.arange(0, attn.resolution[0], step=2, device=device),
+                                    torch.arange(0, attn.resolution[1], step=2, device=device),
+                                    indexing="ij",
+                                )
+                            ).flatten(1)
+                            rel_pos = (q_pos[..., :, None] - k_pos[..., None, :]).abs()
+                            rel_pos = (rel_pos[0] * attn.resolution[1]) + rel_pos[1]
+                            attn.attention_bias_idxs = nn.Buffer(rel_pos.to(torch.long), persistent=False)
 
                     old_base = (old_base[0] // 2, old_base[1] // 2)
                     new_base = (new_base[0] // 2, new_base[1] // 2)
@@ -590,16 +594,22 @@ class EfficientFormer_v2(DetectorBackbone):
                             m.token_mixer.resolution = c_new_base
                             m.token_mixer.N = m.token_mixer.resolution[0] * m.token_mixer.resolution[1]
 
-                            m.token_mixer.attention_biases = nn.Parameter(
-                                interpolate_attention_bias(m.token_mixer.attention_biases, c_old_base, c_new_base)
-                            )
-
-                            pos = torch.stack(
-                                torch.meshgrid(torch.arange(c_new_base[0]), torch.arange(c_new_base[1]), indexing="ij")
-                            ).flatten(1)
-                            rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
-                            rel_pos = (rel_pos[0] * c_new_base[1]) + rel_pos[1]
-                            m.token_mixer.attention_bias_idxs = nn.Buffer(torch.LongTensor(rel_pos), persistent=False)
+                            with torch.no_grad():
+                                m.token_mixer.attention_biases = nn.Parameter(
+                                    interpolate_attention_bias(m.token_mixer.attention_biases, c_old_base, c_new_base)
+                                )
+
+                                device = m.token_mixer.attention_biases.device
+                                pos = torch.stack(
+                                    torch.meshgrid(
+                                        torch.arange(c_new_base[0], device=device),
+                                        torch.arange(c_new_base[1], device=device),
+                                        indexing="ij",
+                                    )
+                                ).flatten(1)
+                                rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
+                                rel_pos = (rel_pos[0] * c_new_base[1]) + rel_pos[1]
+                                m.token_mixer.attention_bias_idxs = nn.Buffer(rel_pos.to(torch.long), persistent=False)
 
 
 registry.register_model_config(

birder/net/efficientvit_msft.py

@@ -497,14 +497,16 @@ class EfficientViT_MSFT(DetectorBackbone):
 
                                 idxs.append(attention_offsets[offset])
 
-                        m.mixer.m.attn.attention_biases = nn.Parameter(
-                            interpolate_attention_bias(
-                                m.mixer.m.attn.attention_biases, old_window_resolution, window_resolution
+                        with torch.no_grad():
+                            m.mixer.m.attn.attention_biases = nn.Parameter(
+                                interpolate_attention_bias(
+                                    m.mixer.m.attn.attention_biases, old_window_resolution, window_resolution
+                                )
+                            )
+                            device = m.mixer.m.attn.attention_biases.device
+                            m.mixer.m.attn.attention_bias_idxs = nn.Buffer(
+                                torch.tensor(idxs, device=device, dtype=torch.long).view(N, N), persistent=False
                             )
-                        )
-                        m.mixer.m.attn.attention_bias_idxs = nn.Buffer(
-                            torch.LongTensor(idxs).view(N, N), persistent=False
-                        )
 
 
 registry.register_model_config(

birder/net/fastvit.py

@@ -879,6 +879,7 @@ class FastViT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         x = self.forward_features(x)
         return self.features(x)
 
+    @torch.no_grad()  # type: ignore[untyped-decorator]
     def reparameterize_model(self) -> None:
         for module in self.modules():
             if hasattr(module, "reparameterize") is True:

birder/net/flexivit.py

@@ -519,15 +519,16 @@ class FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         else:
             num_prefix_tokens = 0
 
-        self.pos_embedding = nn.Parameter(
-            # On rounding error see: https://github.com/facebookresearch/dino/issues/8
-            adjust_position_embedding(
+        with torch.no_grad():
+            pos_embedding = adjust_position_embedding(
+                # On rounding error see: https://github.com/facebookresearch/dino/issues/8
                 self.pos_embedding,
                 (old_size[0] // self.patch_size, old_size[1] // self.patch_size),
                 (new_size[0] // self.patch_size, new_size[1] // self.patch_size),
                 num_prefix_tokens,
             )
-        )
+
+        self.pos_embedding = nn.Parameter(pos_embedding)
 
     def adjust_patch_size(self, patch_size: int) -> None:
         if self.patch_size == patch_size:

birder/net/hiera.py

@@ -612,23 +612,26 @@ class Hiera(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin):
 
         if self.pos_embed_win is not None:
             global_pos_size = (new_size[0] // 2**4, new_size[1] // 2**4)
-            pos_embed = F.interpolate(
-                self.pos_embed,
-                size=global_pos_size,
-                mode="bicubic",
-                antialias=True,
-            )
+            with torch.no_grad():
+                pos_embed = F.interpolate(
+                    self.pos_embed,
+                    size=global_pos_size,
+                    mode="bicubic",
+                    antialias=True,
+                )
+
             self.pos_embed = nn.Parameter(pos_embed)
 
         else:
-            self.pos_embed = nn.Parameter(
-                adjust_position_embedding(
+            with torch.no_grad():
+                pos_embed = adjust_position_embedding(
                     self.pos_embed,
                     (old_size[0] // self.patch_stride[0], old_size[1] // self.patch_stride[1]),
                     (new_size[0] // self.patch_stride[0], new_size[1] // self.patch_stride[1]),
                     0,
                 )
-            )
+
+            self.pos_embed = nn.Parameter(pos_embed)
 
         # Re-init vars
         self.tokens_spatial_shape = [i // s for i, s in zip(new_size, self.patch_stride)]

birder/net/hornet.py

@@ -332,13 +332,15 @@ class HorNet(DetectorBackbone):
                 for m in module.modules():
                     if isinstance(m, HorBlock):
                         if isinstance(m.gn_conv.dwconv, GlobalLocalFilter):
-                            weight = m.gn_conv.dwconv.complex_weight
-                            weight = F.interpolate(
-                                weight.permute(3, 0, 1, 2),
-                                size=(gn_conv_h[i], gn_conv_w[i]),
-                                mode="bilinear",
-                                align_corners=True,
-                            ).permute(1, 2, 3, 0)
+                            with torch.no_grad():
+                                weight = m.gn_conv.dwconv.complex_weight
+                                weight = F.interpolate(
+                                    weight.permute(3, 0, 1, 2),
+                                    size=(gn_conv_h[i], gn_conv_w[i]),
+                                    mode="bilinear",
+                                    align_corners=True,
+                                ).permute(1, 2, 3, 0)
+
                             m.gn_conv.dwconv.complex_weight = nn.Parameter(weight)
 
 

birder/net/iformer.py

@@ -477,12 +477,14 @@ class iFormer(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         resolution = (new_size[0] // 4, new_size[1] // 4)
         for stage in self.body.modules():
             if isinstance(stage, InceptionTransformerStage):
-                orig_dtype = stage.pos_embed.dtype
-                pos_embedding = stage.pos_embed.float()
-                pos_embedding = F.interpolate(
-                    pos_embedding.permute(0, 3, 1, 2), size=resolution, mode="bilinear"
-                ).permute(0, 2, 3, 1)
-                pos_embedding = pos_embedding.to(orig_dtype)
+                with torch.no_grad():
+                    orig_dtype = stage.pos_embed.dtype
+                    pos_embedding = stage.pos_embed.float()
+                    pos_embedding = F.interpolate(
+                        pos_embedding.permute(0, 3, 1, 2), size=resolution, mode="bilinear"
+                    ).permute(0, 2, 3, 1)
+                    pos_embedding = pos_embedding.to(orig_dtype)
+
                 stage.pos_embed = nn.Parameter(pos_embedding)
                 stage.resolution = resolution
                 resolution = (resolution[0] // 2, resolution[1] // 2)

birder/net/levit.py

@@ -454,42 +454,54 @@ class LeViT(BaseNet):
                         # Update Subsample resolution
                         m.q[0].resolution = resolution
 
-                        # Interpolate attention biases
-                        m.attention_biases = nn.Parameter(
-                            interpolate_attention_bias(m.attention_biases, old_resolution, resolution)
-                        )
-
-                        # Rebuild attention bias indices
-                        k_pos = torch.stack(
-                            torch.meshgrid(torch.arange(resolution[0]), torch.arange(resolution[1]), indexing="ij")
-                        ).flatten(1)
-                        q_pos = torch.stack(
-                            torch.meshgrid(
-                                torch.arange(0, resolution[0], step=m.stride),
-                                torch.arange(0, resolution[1], step=m.stride),
-                                indexing="ij",
+                        with torch.no_grad():
+                            # Interpolate attention biases
+                            m.attention_biases = nn.Parameter(
+                                interpolate_attention_bias(m.attention_biases, old_resolution, resolution)
                             )
-                        ).flatten(1)
-                        rel_pos = (q_pos[..., :, None] - k_pos[..., None, :]).abs()
-                        rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
-                        m.attention_bias_idxs = nn.Buffer(rel_pos, persistent=False)
+
+                            # Rebuild attention bias indices
+                            device = m.attention_biases.device
+                            k_pos = torch.stack(
+                                torch.meshgrid(
+                                    torch.arange(resolution[0], device=device),
+                                    torch.arange(resolution[1], device=device),
+                                    indexing="ij",
+                                )
+                            ).flatten(1)
+                            q_pos = torch.stack(
+                                torch.meshgrid(
+                                    torch.arange(0, resolution[0], step=m.stride, device=device),
+                                    torch.arange(0, resolution[1], step=m.stride, device=device),
+                                    indexing="ij",
+                                )
+                            ).flatten(1)
+                            rel_pos = (q_pos[..., :, None] - k_pos[..., None, :]).abs()
+                            rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
+                            m.attention_bias_idxs = nn.Buffer(rel_pos, persistent=False)
 
                         old_resolution = ((old_resolution[0] - 1) // 2 + 1, (old_resolution[1] - 1) // 2 + 1)
                         resolution = ((resolution[0] - 1) // 2 + 1, (resolution[1] - 1) // 2 + 1)
 
                     elif isinstance(m, Attention):
-                        # Interpolate attention biases
-                        m.attention_biases = nn.Parameter(
-                            interpolate_attention_bias(m.attention_biases, old_resolution, resolution)
-                        )
-
-                        # Rebuild attention bias indices
-                        pos = torch.stack(
-                            torch.meshgrid(torch.arange(resolution[0]), torch.arange(resolution[1]), indexing="ij")
-                        ).flatten(1)
-                        rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
-                        rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
-                        m.attention_bias_idxs = nn.Buffer(rel_pos, persistent=False)
+                        with torch.no_grad():
+                            # Interpolate attention biases
+                            m.attention_biases = nn.Parameter(
+                                interpolate_attention_bias(m.attention_biases, old_resolution, resolution)
+                            )
+
+                            # Rebuild attention bias indices
+                            device = m.attention_biases.device
+                            pos = torch.stack(
+                                torch.meshgrid(
+                                    torch.arange(resolution[0], device=device),
+                                    torch.arange(resolution[1], device=device),
+                                    indexing="ij",
+                                )
+                            ).flatten(1)
+                            rel_pos = (pos[..., :, None] - pos[..., None, :]).abs()
+                            rel_pos = (rel_pos[0] * resolution[1]) + rel_pos[1]
+                            m.attention_bias_idxs = nn.Buffer(rel_pos, persistent=False)
 
 
 registry.register_model_config(

birder/net/lit_v1_tiny.py

@@ -340,3 +340,18 @@ registry.register_model_config(
         "drop_path_rate": 0.1,
     },
 )
+
+registry.register_weights(
+    "lit_v1_t_il-common",
+    {
+        "description": "LIT v1 Tiny model trained on the il-common dataset",
+        "resolution": (256, 256),
+        "formats": {
+            "pt": {
+                "file_size": 75.2,
+                "sha256": "93813b2716eb9f33e06dc15ab2ba335c6d219354d2983bbc4f834f8f4e688e5c",
+            }
+        },
+        "net": {"network": "lit_v1_t", "tag": "il-common"},
+    },
+)

birder/net/maxvit.py

@@ -52,8 +52,10 @@ def _make_block_input_shapes(input_size: tuple[int, int], n_blocks: int) -> list
     return shapes
 
 
-def _get_relative_position_index(height: int, width: int) -> torch.Tensor:
-    coords = torch.stack(torch.meshgrid([torch.arange(height), torch.arange(width)], indexing="ij"))
+def _get_relative_position_index(height: int, width: int, device: torch.device | None = None) -> torch.Tensor:
+    coords = torch.stack(
+        torch.meshgrid([torch.arange(height, device=device), torch.arange(width, device=device)], indexing="ij")
+    )
     coords_flat = torch.flatten(coords, 1)
     relative_coords = coords_flat[:, :, None] - coords_flat[:, None, :]
     relative_coords = relative_coords.permute(1, 2, 0).contiguous()
@@ -152,7 +154,9 @@ class RelativePositionalMultiHeadAttention(nn.Module):
         self.relative_position_bias_table = nn.Parameter(
             torch.empty(((2 * self.size[0] - 1) * (2 * self.size[1] - 1), self.n_heads), dtype=torch.float32),
         )
-        self.relative_position_index = nn.Buffer(_get_relative_position_index(self.size[0], self.size[1]))
+        self.relative_position_index = nn.Buffer(
+            _get_relative_position_index(self.size[0], self.size[1], device=self.relative_position_bias_table.device)
+        )
 
         # Initialize with truncated normal the bias
         nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02)
@@ -682,60 +686,68 @@ class MaxViT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
                         old_attn_size = attn.size
                         attn.size = self.partition_size
                         attn.max_seq_len = self.partition_size[0] * self.partition_size[1]
-                        attn.relative_position_index = nn.Buffer(
-                            _get_relative_position_index(attn.size[0], attn.size[1])
-                        )
-
-                        # Interpolate relative_position_bias_table, adapted from
-                        # https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/pos_embed_rel.py
-                        dst_size = (2 * attn.size[0] - 1, 2 * attn.size[1] - 1)
-                        rel_pos_bias = attn.relative_position_bias_table
-                        rel_pos_bias = rel_pos_bias.detach()
-
-                        num_attn_heads = rel_pos_bias.size(1)
-                        src_size = (2 * old_attn_size[0] - 1, 2 * old_attn_size[1] - 1)
-
-                        def _calc(src: int, dst: int) -> list[float]:
-                            (left, right) = 1.01, 1.5
-                            while right - left > 1e-6:
-                                q = (left + right) / 2.0
-                                gp = (1.0 - q ** (src // 2)) / (1.0 - q)  # Geometric progression
-                                if gp > dst // 2:
-                                    right = q
-
-                                else:
-                                    left = q
-
-                            dis = []
-                            cur = 1.0
-                            for i in range(src // 2):
-                                dis.append(cur)
-                                cur += q ** (i + 1)
-
-                            r_ids = [-_ for _ in reversed(dis)]
-                            return r_ids + [0] + dis
-
-                        y = _calc(src_size[0], dst_size[0])
-                        x = _calc(src_size[1], dst_size[1])
-
-                        ty = dst_size[0] // 2.0
-                        tx = dst_size[1] // 2.0
-                        dy = torch.arange(-ty, ty + 0.1, 1.0)
-                        dx = torch.arange(-tx, tx + 0.1, 1.0)
-                        dxy = torch.meshgrid(dx, dy, indexing="ij")
-
-                        all_rel_pos_bias = []
-                        for i in range(num_attn_heads):
-                            z = rel_pos_bias[:, i].view(src_size[0], src_size[1]).float()
-                            rgi = interpolate.RegularGridInterpolator(
-                                (x, y), z.numpy().T, method="cubic", bounds_error=False, fill_value=None
+                        with torch.no_grad():
+                            attn.relative_position_index = nn.Buffer(
+                                _get_relative_position_index(
+                                    attn.size[0],
+                                    attn.size[1],
+                                    device=attn.relative_position_bias_table.device,
+                                )
                             )
-                            r = torch.Tensor(rgi(dxy)).T.contiguous().to(rel_pos_bias.device)
-
-                            r = r.view(-1, 1)
-                            all_rel_pos_bias.append(r)
 
-                        rel_pos_bias = torch.concat(all_rel_pos_bias, dim=-1)
+                            # Interpolate relative_position_bias_table, adapted from
+                            # https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/pos_embed_rel.py
+                            dst_size = (2 * attn.size[0] - 1, 2 * attn.size[1] - 1)
+                            rel_pos_bias = attn.relative_position_bias_table.detach()
+                            rel_pos_device = rel_pos_bias.device
+                            rel_pos_bias = rel_pos_bias.float().cpu()
+
+                            num_attn_heads = rel_pos_bias.size(1)
+                            src_size = (2 * old_attn_size[0] - 1, 2 * old_attn_size[1] - 1)
+
+                            def _calc(src: int, dst: int) -> list[float]:
+                                (left, right) = 1.01, 1.5
+                                while right - left > 1e-6:
+                                    q = (left + right) / 2.0
+                                    gp = (1.0 - q ** (src // 2)) / (1.0 - q)  # Geometric progression
+                                    if gp > dst // 2:
+                                        right = q
+
+                                    else:
+                                        left = q
+
+                                dis = []
+                                cur = 1.0
+                                for i in range(src // 2):
+                                    dis.append(cur)
+                                    cur += q ** (i + 1)
+
+                                r_ids = [-_ for _ in reversed(dis)]
+                                return r_ids + [0] + dis
+
+                            y = _calc(src_size[0], dst_size[0])
+                            x = _calc(src_size[1], dst_size[1])
+
+                            ty = dst_size[0] // 2.0
+                            tx = dst_size[1] // 2.0
+                            dy = torch.arange(-ty, ty + 0.1, 1.0)
+                            dx = torch.arange(-tx, tx + 0.1, 1.0)
+                            dxy = torch.meshgrid(dx, dy, indexing="ij")
+
+                            all_rel_pos_bias = []
+                            for i in range(num_attn_heads):
+                                z = rel_pos_bias[:, i].view(src_size[0], src_size[1])
+                                rgi = interpolate.RegularGridInterpolator(
+                                    (x, y), z.numpy().T, method="cubic", bounds_error=False, fill_value=None
+                                )
+                                r = torch.tensor(
+                                    rgi(dxy), device=rel_pos_device, dtype=rel_pos_bias.dtype
+                                ).T.contiguous()
+
+                                r = r.view(-1, 1)
+                                all_rel_pos_bias.append(r)
+
+                            rel_pos_bias = torch.concat(all_rel_pos_bias, dim=-1)
                         attn.relative_position_bias_table = nn.Parameter(rel_pos_bias)
 
                 new_grid_size = m.grid_size

birder/net/mobileone.py

@@ -380,6 +380,7 @@ class MobileOne(DetectorBackbone):
         x = self.forward_features(x)
         return self.features(x)
 
+    @torch.no_grad()  # type: ignore[untyped-decorator]
     def reparameterize_model(self) -> None:
         for module in self.modules():
             if hasattr(module, "reparameterize") is True:

birder/net/mvit_v2.py

@@ -638,18 +638,19 @@ class MViT_v2(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
                         rel_sp_dim_h = 2 * max(q_size_h, kv_size_h) - 1
                         rel_sp_dim_w = 2 * max(q_size_w, kv_size_w) - 1
 
-                        rel_pos_h = m.attn.rel_pos_h
-                        rel_pos_h_resized = F.interpolate(
-                            rel_pos_h.reshape(1, rel_pos_h.shape[0], -1).permute(0, 2, 1),
-                            size=rel_sp_dim_h,
-                            mode="linear",
-                        )
-                        rel_pos_w = m.attn.rel_pos_w
-                        rel_pos_w_resized = F.interpolate(
-                            rel_pos_w.reshape(1, rel_pos_w.shape[0], -1).permute(0, 2, 1),
-                            size=rel_sp_dim_w,
-                            mode="linear",
-                        )
+                        with torch.no_grad():
+                            rel_pos_h = m.attn.rel_pos_h
+                            rel_pos_h_resized = F.interpolate(
+                                rel_pos_h.reshape(1, rel_pos_h.shape[0], -1).permute(0, 2, 1),
+                                size=rel_sp_dim_h,
+                                mode="linear",
+                            )
+                            rel_pos_w = m.attn.rel_pos_w
+                            rel_pos_w_resized = F.interpolate(
+                                rel_pos_w.reshape(1, rel_pos_w.shape[0], -1).permute(0, 2, 1),
+                                size=rel_sp_dim_w,
+                                mode="linear",
+                            )
 
                         m.attn.rel_pos_h = nn.Parameter(rel_pos_h_resized.reshape(-1, rel_sp_dim_h).permute(1, 0))
                         m.attn.rel_pos_w = nn.Parameter(rel_pos_w_resized.reshape(-1, rel_sp_dim_w).permute(1, 0))

birder/net/pit.py

@@ -258,9 +258,10 @@ class PiT(DetectorBackbone):
         height = (new_size[0] - self.patch_size[0]) // self.patch_stride[0] + 1
         width = (new_size[1] - self.patch_size[1]) // self.patch_stride[1] + 1
 
-        self.pos_embed = nn.Parameter(
-            F.interpolate(self.pos_embed, (height, width), mode="bicubic"), requires_grad=True
-        )
+        with torch.no_grad():
+            pos_embed = F.interpolate(self.pos_embed, (height, width), mode="bicubic")
+
+        self.pos_embed = nn.Parameter(pos_embed)
 
 
 registry.register_model_config(

birder/net/pvt_v1.py

@@ -308,7 +308,10 @@ class PVT_v1(DetectorBackbone):
         s = (new_size[0] // 4, new_size[1] // 4)
         for m in self.body.modules():
             if isinstance(m, PyramidVisionTransformerStage):
-                m.pos_embed = nn.Parameter(adjust_position_embedding(m.pos_embed, old_s, s, 0))
+                with torch.no_grad():
+                    pos_embed = adjust_position_embedding(m.pos_embed, old_s, s, 0)
+
+                m.pos_embed = nn.Parameter(pos_embed)
                 old_s = (old_s[0] // 2, old_s[1] // 2)
                 s = (s[0] // 2, s[1] // 2)
 

birder/net/repghost.py

@@ -338,6 +338,7 @@ class RepGhost(DetectorBackbone):
         x = self.forward_features(x)
         return self.features(x)
 
+    @torch.no_grad()  # type: ignore[untyped-decorator]
     def reparameterize_model(self) -> None:
         for module in self.modules():
             if hasattr(module, "reparameterize") is True: