broccoli-ml 0.33.2__py3-none-any.whl → 0.35.0__py3-none-any.whl

broccoli/transformer.py

@@ -235,25 +235,26 @@ class FeedforwardBlock(nn.Module):
         activation=nn.ReLU,
         activation_kwargs=None,
         dropout=0.0,
-        linear_module=nn.Linear,
+        linear_module_up=nn.Linear,
+        linear_module_down=nn.Linear,
         pre_norm=True,
         normformer=False,
-        raw_input=False,
+        post_norm=True,
+        residual_path=True,
     ):
         super().__init__()
 
+        self.residual_path = residual_path
+        self.post_norm = post_norm
+
+        if self.post_norm:
+            self.layernorm = nn.LayerNorm(output_features)
+
         if activation_kwargs is not None:
             self.activation = activation(**activation_kwargs)
         else:
             self.activation = activation()
 
-        if raw_input:
-            self.memory_type = AnchoredLinear
-            self.memory_bias = False
-        else:
-            self.memory_type = nn.Linear
-            self.memory_bias = True
-
         self.dropout = nn.Dropout(dropout)
 
         self.max_features = (
@@ -265,18 +266,21 @@ class FeedforwardBlock(nn.Module):
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
-                linear_module(input_features, self.max_features),
+                linear_module_up(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                self.memory_type(
-                    ratio * output_features, output_features, bias=self.memory_bias
-                ),
+                linear_module_down(ratio * output_features, output_features),
                 self.dropout,
             ]
         )
 
     def forward(self, x):
-        return self.process(x)
+        if self.residual_path and self.post_norm:
+            return self.layernorm(x + self.process(x))
+        elif self.residual_path:
+            return x + self.process(x)
+        else:
+            return x
 
 
 class TransformerBlock(nn.Module):
@@ -305,11 +309,14 @@ class TransformerBlock(nn.Module):
         causal=False,
         linear_module=nn.Linear,
         pre_norm=True,
+        post_norm=False,
         normformer=False,
     ):
         super().__init__()
 
         self.pre_norm = pre_norm
+        self.post_norm = post_norm
+        self.normformer = normformer
 
         self.identity_probability = identity_probability
 
@@ -340,7 +347,7 @@ class TransformerBlock(nn.Module):
             bos_tokens=bos_tokens,
         )
 
-        # Submodules for the feedforward process
+        # Submodule for the feedforward process
         self.ff = FeedforwardBlock(
             d_model,
             mlp_ratio,
@@ -348,9 +355,12 @@ class TransformerBlock(nn.Module):
             activation=activation,
             activation_kwargs=activation_kwargs,
             dropout=mlp_dropout,
-            linear_module=linear_module,
+            linear_module_up=linear_module,
+            linear_module_down=linear_module,
             pre_norm=pre_norm,
             normformer=normformer,
+            post_norm=post_norm,
+            residual_path=True,
         )
 
     @property
@@ -371,19 +381,18 @@ class TransformerBlock(nn.Module):
         identity_x = shuffled[:identity_count, :, :]
         process_x = shuffled[identity_count:, :, :]
 
+        residual_x = process_x
+
         if self.pre_norm:
-            norm_process_x = self.layer_norm_1(process_x)
-            process_x = process_x + self.attn(
-                norm_process_x, norm_process_x, norm_process_x
-            )
-            process_x = process_x + self.ff(process_x)
-        else:  # post-norm
-            process_x = process_x + self.attn(process_x, process_x, process_x)
             process_x = self.layer_norm_1(process_x)
-            process_x = process_x + self.ff(process_x)
+
+        process_x = residual_x + self.attn(process_x, process_x, process_x)
+
+        if self.post_norm:
             process_x = self.layer_norm_2(process_x)
 
-        # Always post norm as eventually we reach the classification head!
+        process_x = self.ff(process_x)
+
         x = torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
 
         return x
@@ -414,6 +423,7 @@ class TransformerEncoder(nn.Module):
         bos_tokens=0,
         return_bos_tokens=False,
         pre_norm=True,
+        post_norm=False,
         normformer=False,
     ):
         if position_embedding_type == "relative":
@@ -474,6 +484,7 @@ class TransformerEncoder(nn.Module):
                     causal=causal,
                     linear_module=linear_module,
                     pre_norm=pre_norm,
+                    post_norm=post_norm,
                     normformer=normformer,
                 )
                 for i in range(n_layers)

broccoli/vit.py

@@ -53,9 +53,7 @@ class ClassificationHead(nn.Module):
     A general classification head for a ViT
     """
 
-    def __init__(
-        self, d_model, linear_module, n_classes, layer_norm=True, batch_norm=True
-    ):
+    def __init__(self, d_model, linear_module, n_classes, batch_norm=True):
         super().__init__()
         self.d_model = d_model
         self.summarize = GetCLSToken()
@@ -67,7 +65,6 @@ class ClassificationHead(nn.Module):
 
         self.classification_process = nn.Sequential(
             *[
-                nn.LayerNorm(d_model) if layer_norm else nn.Identity(),
                 self.summarize,
                 self.projection,
                 self.batch_norm,
@@ -120,8 +117,12 @@ class ViTEncoder(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         transformer_feedforward_first=True,
+        transformer_initial_ff_residual_path=True,
+        transformer_initial_ff_linear_module_up=None,
+        transformer_initial_ff_linear_module_down=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
+        transformer_post_norm=False,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -296,6 +297,7 @@ class ViTEncoder(nn.Module):
                 return_bos_tokens=transformer_return_bos_tokens,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
+                post_norm=transformer_post_norm,
             )
         else:
             self.transformer = nn.Identity()
@@ -308,10 +310,20 @@ class ViTEncoder(nn.Module):
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
-                linear_module=linear_module,
+                linear_module_up=(
+                    transformer_initial_ff_linear_module_up
+                    if transformer_initial_ff_linear_module_up is not None
+                    else linear_module
+                ),
+                linear_module_down=(
+                    transformer_initial_ff_linear_module_down
+                    if transformer_initial_ff_linear_module_down is not None
+                    else linear_module
+                ),
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
-                raw_input=not cnn,
+                post_norm=transformer_post_norm,
+                residual_path=transformer_initial_ff_residual_path,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -365,8 +377,12 @@ class ViT(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         transformer_feedforward_first=True,
+        transformer_initial_ff_residual_path=True,
+        transformer_initial_ff_linear_module_up=None,
+        transformer_initial_ff_linear_module_down=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
+        transformer_post_norm=False,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -421,8 +437,12 @@ class ViT(nn.Module):
             pooling_kernel_stride=pooling_kernel_stride,
             pooling_padding=pooling_padding,
             transformer_feedforward_first=transformer_feedforward_first,
+            transformer_initial_ff_residual_path=transformer_initial_ff_residual_path,
+            transformer_initial_ff_linear_module_up=transformer_initial_ff_linear_module_up,
+            transformer_initial_ff_linear_module_down=transformer_initial_ff_linear_module_down,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
+            transformer_post_norm=transformer_post_norm,
             transformer_position_embedding=transformer_position_embedding,
             transformer_embedding_size=transformer_embedding_size,
             transformer_layers=transformer_layers,

{broccoli_ml-0.33.2.dist-info → broccoli_ml-0.35.0.dist-info}/METADATA

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

{broccoli_ml-0.33.2.dist-info → broccoli_ml-0.35.0.dist-info}/RECORD

@@ -8,10 +8,10 @@ broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=8Y9vD85ZEgNZsIQgO3uRQ3lOQR-JjwvabY8liCrfNCk,4831
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=ks2TRCdS10k2XvxEieh2sj_LzjTNRuiO6gekKFTtziI,4533
-broccoli/transformer.py,sha256=xXc9dnceGPCOaloITvspNxrkusdSCE-nRn5Xx7-L_XM,17061
+broccoli/transformer.py,sha256=t0gsADJC9UOlwjm7tDKdy0pAZ8l3clTcCnes86zvH-k,17203
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=7lvP2Sak7N5xkKzgzBbwxIrKK2qMwj7GbBIYuoJNxIU,16214
-broccoli_ml-0.33.2.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.33.2.dist-info/METADATA,sha256=1jnADGHLdddLy4GJnHxjX6noNfssrNuQIiZcJQ68ggA,1257
-broccoli_ml-0.33.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.33.2.dist-info/RECORD,,
+broccoli/vit.py,sha256=c-ZRHiLDOoQDJO9OJ51zD9HqaluG33flIwTXQQfms-g,17389
+broccoli_ml-0.35.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.35.0.dist-info/METADATA,sha256=v0JSpcubSGwxA5dFPbDwz2r2oGZWSeqYND1Mu8WOiJY,1257
+broccoli_ml-0.35.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.35.0.dist-info/RECORD,,