broccoli-ml 0.33.2__tar.gz → 0.34.0__tar.gz

{broccoli_ml-0.33.2 → broccoli_ml-0.34.0}/PKG-INFO

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

{broccoli_ml-0.33.2 → broccoli_ml-0.34.0}/broccoli/transformer.py

@@ -238,22 +238,22 @@ class FeedforwardBlock(nn.Module):
         linear_module=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 = (
@@ -268,15 +268,18 @@ class FeedforwardBlock(nn.Module):
                 linear_module(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(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 +308,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
 
@@ -351,6 +357,8 @@ class TransformerBlock(nn.Module):
             linear_module=linear_module,
             pre_norm=pre_norm,
             normformer=normformer,
+            post_norm=post_norm,
+            residual_path=True,
         )
 
     @property
@@ -371,19 +379,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 +421,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 +482,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_ml-0.33.2 → broccoli_ml-0.34.0}/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,11 @@ 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=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,
@@ -308,10 +308,14 @@ class ViTEncoder(nn.Module):
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
-                linear_module=linear_module,
+                linear_module=(
+                    transformer_initial_ff_linear_module
+                    if transformer_initial_ff_linear_module is not None
+                    else linear_module
+                ),
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
-                raw_input=not cnn,
+                residual_path=transformer_initial_ff_residual_path,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -365,8 +369,11 @@ 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=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 +428,11 @@ 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=transformer_initial_ff_linear_module,
             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 → broccoli_ml-0.34.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "0.33.2"
+version = "0.34.0"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}