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

birder/net/regionvit.py

@@ -464,14 +464,14 @@ class RegionViT(DetectorBackbone):
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
         if unfreeze_features is True:
             for param in self.norm.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
         o_x = x
@@ -488,16 +488,16 @@ class RegionViT(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.patch_embed.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
         for param in self.cls_token.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         o_x = x

birder/net/regnet.py

@@ -364,14 +364,14 @@ class RegNet(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_retention(
         self,

birder/net/regnet_z.py

@@ -210,14 +210,14 @@ class RegNet_Z(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_retention(
         self,

birder/net/repghost.py

@@ -321,14 +321,14 @@ class RepGhost(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/repvgg.py

@@ -302,14 +302,14 @@ class RepVgg(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/repvit.py

@@ -399,18 +399,18 @@ class RepViT(DetectorBackbone):
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
             for param in self.dist_classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
         if unfreeze_features is True:
             for param in self.features.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
     def transform_to_backbone(self) -> None:
         self.features = nn.Identity()
@@ -430,14 +430,14 @@ class RepViT(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/resnest.py

@@ -271,14 +271,14 @@ class ResNeSt(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/resnet_v1.py

@@ -192,14 +192,14 @@ class ResNet_v1(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/resnet_v2.py

@@ -178,14 +178,14 @@ class ResNet_v2(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/resnext.py

@@ -216,14 +216,14 @@ class ResNeXt(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/rope_deit3.py

@@ -245,16 +245,16 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
-        self.pos_embedding.requires_grad = False
+        self.pos_embedding.requires_grad_(False)
 
         for idx, module in enumerate(self.encoder.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)

birder/net/rope_flexivit.py

@@ -69,6 +69,8 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", None)
         pre_norm: bool = self.config.get("pre_norm", False)
         post_norm: bool = self.config.get("post_norm", True)
+        qkv_bias: bool = self.config.get("qkv_bias", True)
+        qk_norm: bool = self.config.get("qk_norm", False)
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
         class_token: bool = self.config.get("class_token", True)
         attn_pool_head: bool = self.config.get("attn_pool_head", False)
@@ -118,6 +120,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         self.num_reg_tokens = num_reg_tokens
         self.attn_pool_special_tokens = attn_pool_special_tokens
         self.norm_layer = norm_layer
+        self.norm_layer_eps = norm_layer_eps
         self.mlp_layer = mlp_layer
         self.act_layer = act_layer
         self.rope_rot_type = rope_rot_type
@@ -190,6 +193,8 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             attention_dropout,
             dpr,
             pre_norm=pre_norm,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
             activation_layer=act_layer,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
@@ -231,6 +236,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             rope_temperature=rope_temperature,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             mlp_layer=mlp_layer,
             rope_rot_type=rope_rot_type,
         )
@@ -285,16 +291,16 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
         if unfreeze_features is True:
             if self.attn_pool is not None:
                 for param in self.attn_pool.parameters():
-                    param.requires_grad = True
+                    param.requires_grad_(True)
 
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
@@ -332,16 +338,16 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
-        self.pos_embedding.requires_grad = False
+        self.pos_embedding.requires_grad_(False)
 
         for idx, module in enumerate(self.encoder.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     # pylint: disable=too-many-branches
     def masked_encoding_omission(
@@ -588,6 +594,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             rope_temperature=self.rope_temperature,
             layer_scale_init_value=self.layer_scale_init_value,
             norm_layer=self.norm_layer,
+            norm_layer_eps=self.norm_layer_eps,
             mlp_layer=self.mlp_layer,
             rope_rot_type=self.rope_rot_type,
         )

birder/net/rope_vit.py

@@ -150,6 +150,10 @@ class RoPEAttention(nn.Module):
         attn_drop: float,
         proj_drop: float,
         num_special_tokens: int,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        norm_layer_eps: float = 1e-6,
         rope_rot_type: str = "standard",
     ) -> None:
         super().__init__()
@@ -167,7 +171,14 @@ class RoPEAttention(nn.Module):
         else:
             raise ValueError(f"Unknown rope_rot_type, got '{rope_rot_type}'")
 
-        self.qkv = nn.Linear(dim, dim * 3)
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        if qk_norm is True:
+            self.q_norm = norm_layer(self.head_dim, eps=norm_layer_eps)
+            self.k_norm = norm_layer(self.head_dim, eps=norm_layer_eps)
+        else:
+            self.q_norm = nn.Identity()
+            self.k_norm = nn.Identity()
+
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
@@ -176,6 +187,8 @@ class RoPEAttention(nn.Module):
         (B, N, C) = x.size()
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
         (q, k, v) = qkv.unbind(0)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
 
         n = self.num_special_tokens
         q = torch.concat([q[:, :, :n, :], self.apply_rot_fn(q[:, :, n:, :], rope)], dim=2)
@@ -207,6 +220,8 @@ class EncoderBlock(nn.Module):
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
         norm_layer_eps: float = 1e-6,
         mlp_layer: Callable[..., nn.Module] = FFN,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
         rope_rot_type: str = "standard",
     ) -> None:
         super().__init__()
@@ -222,6 +237,10 @@ class EncoderBlock(nn.Module):
             attn_drop=attention_dropout,
             proj_drop=dropout,
             num_special_tokens=num_special_tokens,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             rope_rot_type=rope_rot_type,
         )
         if layer_scale_init_value is not None:
@@ -249,7 +268,6 @@ class EncoderBlock(nn.Module):
 
 
 class Encoder(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         num_layers: int,
@@ -261,6 +279,8 @@ class Encoder(nn.Module):
         attention_dropout: float,
         dpr: list[float],
         pre_norm: bool = False,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
         activation_layer: Callable[..., nn.Module] = nn.GELU,
         layer_scale_init_value: Optional[float] = None,
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
@@ -293,6 +313,8 @@ class Encoder(nn.Module):
                     norm_layer=norm_layer,
                     norm_layer_eps=norm_layer_eps,
                     mlp_layer=mlp_layer,
+                    qkv_bias=qkv_bias,
+                    qk_norm=qk_norm,
                     rope_rot_type=rope_rot_type,
                 )
             )
@@ -331,6 +353,7 @@ class MAEDecoderBlock(nn.Module):
         rope_temperature: float,
         layer_scale_init_value: Optional[float] = None,
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        norm_layer_eps: float = 1e-6,
         mlp_layer: Callable[..., nn.Module] = FFN,
         rope_rot_type: str = "standard",
     ) -> None:
@@ -346,7 +369,7 @@ class MAEDecoderBlock(nn.Module):
         )
 
         # Attention block
-        self.norm1 = norm_layer(hidden_dim, eps=1e-6)
+        self.norm1 = norm_layer(hidden_dim, eps=norm_layer_eps)
         self.attn = RoPEAttention(
             hidden_dim,
             num_heads,
@@ -361,7 +384,7 @@ class MAEDecoderBlock(nn.Module):
             self.layer_scale_1 = nn.Identity()
 
         # MLP block
-        self.norm2 = norm_layer(hidden_dim, eps=1e-6)
+        self.norm2 = norm_layer(hidden_dim, eps=norm_layer_eps)
         self.mlp = mlp_layer(hidden_dim, mlp_dim, act_layer=activation_layer, dropout=0.0)
         if layer_scale_init_value is not None:
             self.layer_scale_2 = LayerScale(hidden_dim, layer_scale_init_value)
@@ -403,6 +426,8 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", None)
         pre_norm: bool = self.config.get("pre_norm", False)
         post_norm: bool = self.config.get("post_norm", True)
+        qkv_bias: bool = self.config.get("qkv_bias", True)
+        qk_norm: bool = self.config.get("qk_norm", False)
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
         class_token: bool = self.config.get("class_token", True)
         attn_pool_head: bool = self.config.get("attn_pool_head", False)
@@ -450,6 +475,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         self.num_reg_tokens = num_reg_tokens
         self.attn_pool_special_tokens = attn_pool_special_tokens
         self.norm_layer = norm_layer
+        self.norm_layer_eps = norm_layer_eps
         self.mlp_layer = mlp_layer
         self.act_layer = act_layer
         self.rope_rot_type = rope_rot_type
@@ -521,6 +547,8 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             attention_dropout,
             dpr,
             pre_norm=pre_norm,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
             activation_layer=act_layer,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
@@ -562,6 +590,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             rope_temperature=rope_temperature,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             mlp_layer=mlp_layer,
             rope_rot_type=rope_rot_type,
         )
@@ -614,16 +643,16 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
         if unfreeze_features is True:
             if self.attn_pool is not None:
                 for param in self.attn_pool.parameters():
-                    param.requires_grad = True
+                    param.requires_grad_(True)
 
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
@@ -661,16 +690,16 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
-        self.pos_embedding.requires_grad = False
+        self.pos_embedding.requires_grad_(False)
 
         for idx, module in enumerate(self.encoder.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_omission(
         self,
@@ -904,6 +933,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             rope_temperature=self.rope_temperature,
             layer_scale_init_value=self.layer_scale_init_value,
             norm_layer=self.norm_layer,
+            norm_layer_eps=self.norm_layer_eps,
             mlp_layer=self.mlp_layer,
             rope_rot_type=self.rope_rot_type,
         )
@@ -931,6 +961,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
 #     - rms         : RMSNorm (instead of LayerNorm)
 #     - pn          : Pre-Norm (layer norm before the encoder) - implies different norm eps
 #     - npn         : No Post Norm (disables post-normalization layer)
+#     - qkn         : QK Norm
 #
 #     Feed-Forward Network:
 #     - swiglu      : SwiGLU FFN layer type (instead of standard FFN)
@@ -1068,6 +1099,20 @@ registry.register_model_config(
         "drop_path_rate": 0.1,
     },
 )
+registry.register_model_config(
+    "rope_vit_b16_qkn_ls",
+    RoPE_ViT,
+    config={
+        "patch_size": 16,
+        "num_layers": 12,
+        "num_heads": 12,
+        "hidden_dim": 768,
+        "mlp_dim": 3072,
+        "layer_scale_init_value": 1e-5,
+        "qk_norm": True,
+        "drop_path_rate": 0.1,
+    },
+)
 registry.register_model_config(
     "rope_i_vit_b16_pn_aps_c1",  # For PE Core - https://arxiv.org/abs/2504.13181
     RoPE_ViT,
@@ -1310,6 +1355,20 @@ registry.register_model_config(
         "drop_path_rate": 0.0,
     },
 )
+registry.register_model_config(
+    "rope_vit_reg4_m14_avg",
+    RoPE_ViT,
+    config={
+        "patch_size": 14,
+        "num_layers": 12,
+        "num_heads": 8,
+        "hidden_dim": 512,
+        "mlp_dim": 2048,
+        "num_reg_tokens": 4,
+        "class_token": False,
+        "drop_path_rate": 0.0,
+    },
+)
 registry.register_model_config(
     "rope_vit_reg4_b32",
     RoPE_ViT,

birder/net/shufflenet_v1.py

@@ -220,14 +220,14 @@ class ShuffleNet_v1(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/shufflenet_v2.py

@@ -166,14 +166,14 @@ class ShuffleNet_v2(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/smt.py

@@ -275,7 +275,6 @@ class Stem(nn.Module):
 
 
 class SMTStage(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         dim: int,
@@ -429,7 +428,7 @@ class SMT(DetectorBackbone):
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         return self.body(x)

birder/net/squeezenext.py

@@ -177,14 +177,14 @@ class SqueezeNext(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/ssl/byol.py

@@ -82,8 +82,9 @@ class BYOL(SSLBaseNet):
         online_predictions = self.online_predictor(projection)
         (online_pred_one, online_pred_two) = online_predictions.chunk(2, dim=0)
 
-        target_projections = self.target_encoder(x)
-        (target_proj_one, target_proj_two) = target_projections.chunk(2, dim=0)
+        with torch.no_grad():
+            target_projections = self.target_encoder(x)
+            (target_proj_one, target_proj_two) = target_projections.chunk(2, dim=0)
 
         loss_one = loss_fn(online_pred_one, target_proj_two.detach())
         loss_two = loss_fn(online_pred_two, target_proj_one.detach())