broccoli-ml 0.11.0__tar.gz → 0.13.0__tar.gz

{broccoli_ml-0.11.0 → broccoli_ml-0.13.0}/PKG-INFO

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

broccoli_ml-0.13.0/broccoli/linear.py

@@ -0,0 +1,88 @@
+# UNDER CONSTRUCTION
+
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .tensor import SigmaReparamTensor
+
+
+class SpectralNormLinear(nn.Module):
+    """
+    ...
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        self.weight_init = nn.Parameter(torch.empty(out_features, in_features))
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight_init, a=math.sqrt(5))
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = SigmaReparamTensor(self.weight_init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        # Optional: A nice representation for printing the module.
+        return (
+            f"SpectralNormFeedForward(in_features={self.in_features}",
+            f"out_features={self.out_features}, bias={self.use_bias})",
+        )
+
+
+class RandomLinear(nn.Linear):
+    """ """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,  # <---- TODO: explain this
+        beta=0.1,
+        forward_looks_random=True,
+    ):
+        super().__init__(in_features, out_features, bias=False)
+        self.beta = beta
+        self.forward_looks_random = forward_looks_random
+
+    def forward(self, inputs: torch.Tensor):
+        if not self.training:
+            return F.linear(inputs, self.weight)
+        else:
+            # Initialise self.random_weights
+            random_weights = torch.empty_like(self.weight)
+            nn.init.trunc_normal_(random_weights)
+            random_weights *= self.beta
+
+            if self.forward_looks_random:
+                # Forward using a reparameterisation trick
+                a = F.linear(inputs.detach(), self.weight, self.bias)
+                b = F.linear(inputs, random_weights, bias=None)
+            else:
+                # Forward as (W_actual * input + W_random * input) + bias
+                a = F.linear(inputs, self.weight, self.bias)
+                b = F.linear(inputs, random_weights, bias=None)
+
+            return a + b

{broccoli_ml-0.11.0 → broccoli_ml-0.13.0}/broccoli/transformer.py

@@ -10,6 +10,7 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from .rope import RotaryEmbedding, apply_rotary_emb
+from .linear import SpectralNormLinear
 
 
 class MHAttention(nn.Module):
@@ -245,19 +246,19 @@ class FeedforwardLayer(nn.Module):
 
         self.dropout = nn.Dropout(dropout)
 
+        self.max_features = (
+            2 * ratio * output_features
+            if activation.__name__.endswith("GLU")
+            else ratio * output_features
+        )
+
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features),
-                linear_module(
-                    input_features,
-                    (
-                        2 * ratio * output_features
-                        if activation.__name__.endswith("GLU")
-                        else ratio * output_features
-                    ),
-                ),
+                linear_module(input_features, self.max_features),
                 self.activation,
-                linear_module(ratio * output_features, output_features),
+                nn.LayerNorm(self.max_features),
+                linear_module(ratio * output_features, output_features, bias=False),
                 self.dropout,
             ]
         )
@@ -296,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":

{broccoli_ml-0.11.0 → broccoli_ml-0.13.0}/broccoli/vit.py

@@ -20,37 +20,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):
@@ -88,6 +125,7 @@ class ViTEncoder(nn.Module):
         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,
@@ -249,6 +287,7 @@ 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()
@@ -298,7 +337,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)
@@ -332,6 +371,7 @@ class CCT(nn.Module):
         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,
@@ -341,6 +381,7 @@ class CCT(nn.Module):
         initial_batch_norm=True,
         linear_module=nn.Linear,
         image_classes=100,
+        head=SequencePoolClassificationHead,
     ):
 
         super().__init__()
@@ -385,6 +426,7 @@ 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,
@@ -393,7 +435,8 @@ class CCT(nn.Module):
             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.11.0 → broccoli_ml-0.13.0}/pyproject.toml

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

broccoli_ml-0.11.0/broccoli/linear.py

@@ -1,41 +0,0 @@
-# UNDER CONSTRUCTION
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-
-class RandomLinear(nn.Linear):
-    """ """
-
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        bias: bool = False,  # <---- TODO: explain this
-        beta=0.1,
-        forward_looks_random=True,
-    ):
-        super().__init__(in_features, out_features, bias=False)
-        self.beta = beta
-        self.forward_looks_random = forward_looks_random
-
-    def forward(self, inputs: torch.Tensor):
-        if not self.training:
-            return F.linear(inputs, self.weight)
-        else:
-            # Initialise self.random_weights
-            random_weights = torch.empty_like(self.weight)
-            nn.init.trunc_normal_(random_weights)
-            random_weights *= self.beta
-
-            if self.forward_looks_random:
-                # Forward using a reparameterisation trick
-                a = F.linear(inputs.detach(), self.weight, self.bias)
-                b = F.linear(inputs, random_weights, bias=None)
-            else:
-                # Forward as (W_actual * input + W_random * input) + bias
-                a = F.linear(inputs, self.weight, self.bias)
-                b = F.linear(inputs, random_weights, bias=None)
-
-            return a + b