broccoli-ml 11.0.0__py3-none-any.whl → 12.1.0__py3-none-any.whl

broccoli/transformer.py

@@ -21,29 +21,6 @@ except ImportError:
     FLASH_ATTN = False
 
 
-class LayerScale(nn.Module):
-    def __init__(self, dim, decay=False, init_values=1e-4):
-        super().__init__()
-        self.dim = dim
-        self.decay = decay
-        self.init_values = init_values
-        self.reset_parameters()
-
-    def forward(self, x):
-        if self.decay:
-            return x * self.scale
-        else:
-            return x * self.nondecay_scale
-
-    def reset_parameters(self):
-        if self.decay:
-            self.scale = nn.Parameter(self.init_values * torch.ones(self.dim))
-            self.nondecay_scale = None
-        else:
-            self.nondecay_scale = nn.Parameter(self.init_values * torch.ones(self.dim))
-            self.scale = None
-
-
 def drop_path(
     x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
 ):
@@ -527,7 +504,6 @@ class TransformerBlock(nn.Module):
         post_norm=False,
         normformer=False,
         checkpoint_ff=True,
-        layerscale=True,
     ):
         """
         Args:
@@ -556,14 +532,6 @@ class TransformerBlock(nn.Module):
             self.post_attention_norm = nn.LayerNorm(d_model)
             self.post_mlp_norm = nn.LayerNorm(d_model)
 
-        self.layerscale = layerscale
-        if layerscale:
-            self.layerscale1 = LayerScale(d_model)
-            self.layerscale2 = LayerScale(d_model)
-        else:
-            self.layerscale1 = nn.Identity()
-            self.layerscale2 = nn.Identity()
-
         if relative_position_embedding:
             max_freq = int(max(source_size) / 2)  # Suggested by Gemini!
             if d_model < 16:
@@ -624,35 +592,52 @@ class TransformerBlock(nn.Module):
         return self.attn._kv_distance
 
     def forward(self, x):
-
         if self.pre_norm:
-            process_x = self.pre_attention_norm(x)
-        else:
-            process_x = x
+            x = self.layer_norm_1(x)
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = self.layer_norm_2(x)
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
+            if self.post_norm:  # i.e. in addition! Pre and post.
+                x = self.layer_norm_3(x)
+        elif self.post_norm:  # i.e. only, not prenorm, just post
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = self.layer_norm_1(x)
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
+            x = self.layer_norm_2(x)
+        else:  # Not pre or post norm. Stand well back.
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
 
-        processed = self.drop_path(
-            self.layerscale1(self.attn(process_x, process_x, process_x))
-        )
+        # if self.pre_norm:
+        #     process_x = self.pre_attention_norm(x)
+        # else:
+        #     process_x = x
 
-        if self.normformer:
-            processed = self.normformer_norm(processed)
+        # processed = self.drop_path(self.attn(process_x, process_x, process_x))
 
-        x = x + processed
+        # if self.normformer:
+        #     processed = self.normformer_norm(processed)
 
-        if self.post_norm:
-            x = self.post_attention_norm(x)
+        # if self.residual_path:
+        #     x = x + processed
 
-        if self.pre_norm:
-            process_x = self.pre_mlp_norm(x)
-        else:
-            process_x = x
+        # if self.post_norm:
+        #     x = self.post_attention_norm(x)
 
-        x = x + self.drop_path(self.layerscale2(self.ff(process_x)))
+        # if self.pre_norm:
+        #     process_x = self.pre_mlp_norm(x)
+        # else:
+        #     process_x = x
 
-        if self.post_norm:
-            x = self.post_mlp_norm(x)
+        # processed = self.drop_path(self.ff(process_x))
 
-        return x
+        # if self.residual_path:
+        #     x = x + processed
+
+        # if self.post_norm:
+        #     x = self.post_mlp_norm(x)
+
+        # return x
 
     def attention_logits(self, x):
         """
@@ -682,10 +667,6 @@ class TransformerBlock(nn.Module):
         self.attn.reset_parameters()
         self.ff.reset_parameters()
 
-        if self.layerscale:
-            self.layerscale1.reset_parameters()
-            self.layerscale2.reset_parameters()
-
 
 class TransformerEncoder(nn.Module):
     """
@@ -722,7 +703,6 @@ class TransformerEncoder(nn.Module):
         normformer=False,
         msa_scaling="d",
         checkpoint_ff=True,
-        layerscale=True,
     ):
         """
         Args:
@@ -757,12 +737,6 @@ class TransformerEncoder(nn.Module):
         self._utility_tokens = utility_tokens
         self.return_utility_tokens = return_utility_tokens
 
-        if layerscale:
-            rope_and_ape = absolute_position_embedding and relative_position_embedding
-            self.position_layerscale = LayerScale(d_model, decay=rope_and_ape)
-        else:
-            self.position_layerscale = None
-
         # Initialise utility tokens with normal init, like usual Pytorch embeddings
         if self._utility_tokens:
             self._utility_token_embedding = nn.Parameter(
@@ -827,7 +801,6 @@ class TransformerEncoder(nn.Module):
                     post_norm=post_norm,
                     normformer=normformer,
                     checkpoint_ff=checkpoint_ff,
-                    layerscale=layerscale,
                 )
                 for i in range(n_layers)
             ]
@@ -855,8 +828,6 @@ class TransformerEncoder(nn.Module):
                     0
                 )  # to shape (1, seq_len) to broadcast over batch
             )
-            if self.position_layerscale is not None:
-                position_embedding = self.position_layerscale(position_embedding)
             x += position_embedding
 
         return x

broccoli/vit.py

@@ -187,7 +187,6 @@ class ViTEncoder(nn.Module):
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
-        transformer_layerscale=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -353,7 +352,6 @@ class ViTEncoder(nn.Module):
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 checkpoint_ff=transformer_checkpoint_ff,
-                layerscale=transformer_layerscale,
             )
         else:
             self.transformer = nn.Identity()
@@ -489,7 +487,6 @@ class ViT(nn.Module):
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
-        transformer_layerscale=True,
         head=SequencePoolClassificationHead,
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
@@ -562,7 +559,6 @@ class ViT(nn.Module):
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,
-            transformer_layerscale=transformer_layerscale,
             linear_module=linear_module,
         )
 

{broccoli_ml-11.0.0.dist-info → broccoli_ml-12.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 11.0.0
+Version: 12.1.0
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-11.0.0.dist-info → broccoli_ml-12.1.0.dist-info}/RECORD

@@ -4,10 +4,10 @@ broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
 broccoli/linear.py,sha256=W-3aNpBjd_0xRyzbCKkmg4H1qmslQOIQhB-WDDay2nM,13125
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=urW9KGxK47Mns7ZowPKTgdEyp4Yd21uw3PNCpKx04cI,29223
+broccoli/transformer.py,sha256=cPXHGlfkvJAXgr_ONyESQ0RAEzn7yqZEoyZ3l4N3EX8,28571
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=EGbQb-atuzG3JAx7kdTaJEbWvQR-4XgyYvwjKkN5C38,22612
-broccoli_ml-11.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-11.0.0.dist-info/METADATA,sha256=7ixQROHXr3LyCMhT5hkKmvip8C8Bytmtz-IoQCbC_NQ,1369
-broccoli_ml-11.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-11.0.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=DvVpayMIcUhH7Xg6CiYyeedUuuMHrjsGxEdXfnTGa_Q,22428
+broccoli_ml-12.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-12.1.0.dist-info/METADATA,sha256=2j4nt5aHVcPiW_YF2Sx4Yu_l3MKe2fa0S9v3aQ1PZE0,1369
+broccoli_ml-12.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-12.1.0.dist-info/RECORD,,