broccoli-ml 0.12.0__tar.gz → 0.13.1__tar.gz

{broccoli_ml-0.12.0 → broccoli_ml-0.13.1}/PKG-INFO

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

{broccoli_ml-0.12.0 → broccoli_ml-0.13.1}/broccoli/linear.py

@@ -10,7 +10,8 @@ from .tensor import SigmaReparamTensor
 
 class SpectralNormLinear(nn.Module):
     """
-    ...
+    Inspired by Apple's Spectral Normed Linear Layers
+        (https://github.com/apple/ml-sigma-reparam)
     """
 
     def __init__(self, in_features: int, out_features: int, bias: bool = True):

{broccoli_ml-0.12.0 → broccoli_ml-0.13.1}/broccoli/transformer.py

@@ -222,7 +222,7 @@ class MHAttention(nn.Module):
         return self.out_proj(output_without_heads)
 
 
-class FeedforwardLayer(nn.Module):
+class FeedforwardBlock(nn.Module):
     """
     ...
     """
@@ -236,7 +236,6 @@ class FeedforwardLayer(nn.Module):
         activation_kwargs=None,
         dropout=0.0,
         linear_module=nn.Linear,
-        norm_memory=False,
     ):
         super().__init__()
 
@@ -253,22 +252,13 @@ class FeedforwardLayer(nn.Module):
             else ratio * output_features
         )
 
-        if norm_memory:
-            self.memory_type = SpectralNormLinear
-            self.bias_memories = False
-        else:
-            self.memory_type = linear_module
-            self.bias_memories = True
-
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features),
                 linear_module(input_features, self.max_features),
                 self.activation,
                 nn.LayerNorm(self.max_features),
-                self.memory_type(
-                    ratio * output_features, output_features, bias=self.bias_memories
-                ),
+                linear_module(ratio * output_features, output_features, bias=False),
                 self.dropout,
             ]
         )
@@ -307,7 +297,6 @@ class TransformerBlock(nn.Module):
 
         self.identity_probability = identity_probability
 
-        # Submodules for applying attention
         self.layer_norm = nn.LayerNorm(d_model)
 
         if position_embedding_type == "relative":
@@ -335,7 +324,7 @@ class TransformerBlock(nn.Module):
         )
 
         # Submodules for the feedforward process
-        self.ff = FeedforwardLayer(
+        self.ff = FeedforwardBlock(
             d_model,
             mlp_ratio,
             d_model,

{broccoli_ml-0.12.0 → broccoli_ml-0.13.1}/broccoli/vit.py

@@ -1,9 +1,10 @@
 import math
 from typing import Optional
 
-from .transformer import TransformerEncoder, FeedforwardLayer
+from .transformer import TransformerEncoder, FeedforwardBlock
 from .cnn import SpaceToDepth, calculate_output_spatial_size, spatial_tuple
 from .activation import ReLU, SquaredReLU, GELU, SwiGLU
+from .linear import SpectralNormLinear
 from einops import einsum
 from einops.layers.torch import Rearrange
 import torch.nn as nn
@@ -20,37 +21,74 @@ class PadTensor(nn.Module):
         return F.pad(x, *self.args, **self.kwargs)
 
 
+class GetCLSToken(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x[:, 0, :]
+
+
 class SequencePool(nn.Module):
+    def __init__(self, d_model, linear_module):
+        super().__init__()
+        self.attention = nn.Sequential(
+            *[
+                linear_module(d_model, 1),
+                Rearrange("batch seq 1 -> batch seq"),
+                nn.Softmax(dim=-1),
+            ]
+        )
+
+    def forward(self, x):
+        weights = self.attention(x)
+        return einsum(weights, x, "batch seq, batch seq d_model -> batch d_model")
+
+
+class ClassificationHead(nn.Module):
     """
-    As described in [Hasani et al. (2021) *''Escaping the Big Data Paradigm with
-        Compact Transformers''*](https://arxiv.org/abs/2104.05704). It can be viewed
-        as a generalisation of average pooling.
+    A general classification head for a ViT
     """
 
-    def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
+    def __init__(self, d_model, linear_module, n_classes, batch_norm=True):
         super().__init__()
         self.d_model = d_model
-        self.attention = nn.Sequential(
+        self.summarize = GetCLSToken()
+        self.process = nn.Sequential(
             *[
                 linear_module(d_model, 1),
                 Rearrange("batch seq 1 -> batch seq"),
                 nn.Softmax(dim=-1),
             ]
         )
-        self.projection = nn.Linear(d_model, out_dim)
-        self.batch_norm = batch_norm
+        self.projection = nn.Linear(d_model, n_classes)
         if batch_norm:
-            self.norm = nn.BatchNorm1d(out_dim, affine=False)
+            self.batch_norm = nn.BatchNorm1d(n_classes, affine=False)
         else:
-            self.norm = None
+            self.batch_norm = nn.Identity()
 
-    def forward(self, x):
-        weights = self.attention(x)
-        weighted_embedding = einsum(
-            weights, x, "batch seq, batch seq d_model -> batch d_model"
+        self.classification_process = nn.Sequential(
+            *[
+                self.summarize,
+                self.projection,
+                self.batch_norm,
+            ]
         )
-        projection = self.projection(weighted_embedding)
-        return self.norm(projection) if self.batch_norm else projection
+
+    def forward(self, x):
+        return self.classification_process(x)
+
+
+class SequencePoolClassificationHead(ClassificationHead):
+    """
+    As described in [Hasani et al. (2021) *''Escaping the Big Data Paradigm with
+        Compact Transformers''*](https://arxiv.org/abs/2104.05704). It can be viewed
+        as a generalisation of average pooling.
+    """
+
+    def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
+        super().__init__(d_model, linear_module, out_dim, batch_norm=True)
+        self.summarize = SequencePool()
 
 
 class ViTEncoder(nn.Module):
@@ -66,6 +104,7 @@ class ViTEncoder(nn.Module):
     def __init__(
         self,
         input_size=(32, 32),
+        initial_batch_norm=True,
         cnn=True,
         cnn_in_channels=3,
         cnn_out_channels=16,
@@ -82,20 +121,19 @@ class ViTEncoder(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         intermediate_feedforward_layer=True,
-        norm_intermediate_ff_memory=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
         transformer_bos_tokens=0,
+        transformer_return_bos_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         linear_module=nn.Linear,
-        initial_batch_norm=True,
     ):
         super().__init__()
 
@@ -250,13 +288,14 @@ class ViTEncoder(nn.Module):
                 causal=False,
                 linear_module=linear_module,
                 bos_tokens=transformer_bos_tokens,
+                return_bos_tokens=transformer_return_bos_tokens,
             )
         else:
             self.transformer = nn.Identity()
 
         if intermediate_feedforward_layer:
             self.pooling_channels_padding = nn.Identity()
-            self.intermediate_feedforward_layer = FeedforwardLayer(
+            self.intermediate_feedforward_layer = FeedforwardBlock(
                 pooling_out_channels,
                 transformer_mlp_ratio,
                 transformer_embedding_size,
@@ -264,7 +303,6 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
                 linear_module=linear_module,
-                norm_memory=norm_intermediate_ff_memory,
             )
         elif pooling_out_channels < transformer_embedding_size:
             self.intermediate_feedforward_layer = nn.Identity()
@@ -300,7 +338,7 @@ class ViTEncoder(nn.Module):
         return self.encoder(x)
 
 
-class CCT(nn.Module):
+class ViT(nn.Module):
     """
     Denoising convolutional transformer
     Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
@@ -312,6 +350,8 @@ class CCT(nn.Module):
     def __init__(
         self,
         input_size=(32, 32),
+        image_classes=100,
+        initial_batch_norm=True,
         cnn=True,
         cnn_in_channels=3,
         cnn_out_channels=16,
@@ -328,22 +368,21 @@ class CCT(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         intermediate_feedforward_layer=True,
-        norm_intermediate_ff_memory=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
         transformer_bos_tokens=0,
+        transformer_return_bos_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         batch_norm_outputs=True,
-        initial_batch_norm=True,
-        linear_module=nn.Linear,
-        image_classes=100,
+        linear_module=SpectralNormLinear,
+        head=SequencePoolClassificationHead,
     ):
 
         super().__init__()
@@ -366,6 +405,7 @@ class CCT(nn.Module):
 
         self.encoder = ViTEncoder(
             input_size=input_size,
+            initial_batch_norm=initial_batch_norm,
             cnn=cnn,
             cnn_in_channels=cnn_in_channels,
             cnn_out_channels=cnn_out_channels,
@@ -388,15 +428,16 @@ class CCT(nn.Module):
             transformer_heads=transformer_heads,
             transformer_mlp_ratio=transformer_mlp_ratio,
             transformer_bos_tokens=transformer_bos_tokens,
+            transformer_return_bos_tokens=transformer_return_bos_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             linear_module=linear_module,
-            initial_batch_norm=initial_batch_norm,
         )
-        self.pool = SequencePool(
+
+        self.pool = head(
             transformer_embedding_size,
             linear_module,
             image_classes,

{broccoli_ml-0.12.0 → broccoli_ml-0.13.1}/pyproject.toml

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