broccoli-ml 0.29.1__py3-none-any.whl → 10.0.1__py3-none-any.whl

broccoli/utils.py

@@ -1,9 +1,15 @@
-import importlib.resources
-import torch
+import torch.nn as nn
+import torch.nn.functional as F
 
 
-def get_weights(name: str) -> torch.Tensor:
-    resource_path = importlib.resources.files("broccoli.assets") / name
-    with importlib.resources.as_file(resource_path) as path_to_weights:
-        weights = torch.load(path_to_weights)
-        return weights
+class PadTensor(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.kwargs = kwargs
+
+    def forward(self, x):
+        if sum(self.args[0]) == 0:
+            return x
+        else:
+            return F.pad(x, *self.args, **self.kwargs)

broccoli/vit.py

@@ -4,24 +4,13 @@ from typing import Optional
 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 .utils import PadTensor
+
 from einops import einsum
 from einops.layers.torch import Rearrange
-import torch.nn as nn
-import torch.nn.functional as F
 
-
-class PadTensor(nn.Module):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-        self.args = args
-        self.kwargs = kwargs
-
-    def forward(self, x):
-        if sum(self.args[0]) == 0:
-            return x
-        else:
-            return F.pad(x, *self.args, **self.kwargs)
+import torch
+import torch.nn as nn
 
 
 class GetCLSToken(nn.Module):
@@ -43,22 +32,45 @@ class SequencePool(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         weights = self.attention(x)
         return einsum(weights, x, "batch seq, batch seq d_model -> batch d_model")
 
+    def attention_scores(self, x):
+        return self.attention(x)
+
+    def reset_parameters(self):
+        # Iterate over modules in the sequential block
+        for module in self.attention:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class ClassificationHead(nn.Module):
     """
     A general classification head for a ViT
     """
 
-    def __init__(self, d_model, linear_module, n_classes, batch_norm=True):
+    def __init__(
+        self,
+        d_model,
+        n_classes,
+        logit_projection_layer=nn.Linear,
+        batch_norm_logits=True,
+    ):
         super().__init__()
         self.d_model = d_model
         self.summarize = GetCLSToken()
-        self.projection = nn.Linear(d_model, n_classes)
-        if batch_norm:
+
+        if d_model == n_classes:
+            # No need to project
+            self.projection = nn.Identity()
+        else:
+            self.projection = logit_projection_layer(d_model, n_classes)
+
+        if batch_norm_logits:
             self.batch_norm = nn.BatchNorm1d(n_classes, affine=False)
         else:
             self.batch_norm = nn.Identity()
@@ -71,9 +83,16 @@ class ClassificationHead(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         return self.classification_process(x)
 
+    def reset_parameters(self):
+        for module in self.classification_process:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class SequencePoolClassificationHead(ClassificationHead):
     """
@@ -82,9 +101,31 @@ class SequencePoolClassificationHead(ClassificationHead):
         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=batch_norm)
-        self.summarize = SequencePool(d_model, linear_module)
+    def __init__(
+        self,
+        d_model,
+        n_classes,
+        logit_projection_layer=nn.Linear,
+        batch_norm_logits=True,
+    ):
+        super().__init__(
+            d_model,
+            n_classes,
+            logit_projection_layer=logit_projection_layer,
+            batch_norm_logits=batch_norm_logits,
+        )
+
+        self.summarize = SequencePool(d_model, logit_projection_layer)
+        # Rebuild the classification process with the correct summary module:
+        self.classification_process = nn.Sequential(
+            *[
+                self.summarize,
+                self.projection,
+                self.batch_norm,
+            ]
+        )
+
+        self.reset_parameters()
 
 
 class ViTEncoder(nn.Module):
@@ -117,20 +158,36 @@ 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_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
-        transformer_position_embedding="relative",  # absolute or relative
+        transformer_post_norm=False,
+        transformer_absolute_position_embedding=False,
+        transformer_relative_position_embedding=True,
         transformer_embedding_size=256,
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=0,
-        transformer_return_bos_tokens=False,
+        transformer_utility_tokens=0,
+        transformer_talking_heads=False,
+        transformer_return_utility_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
-        transformer_mlp_dropout=0.0,
+        transformer_ff_linear_module_up=None,
+        transformer_ff_linear_module_down=None,
+        transformer_msa_scaling="d",
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
+        transformer_layerscale=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -232,13 +289,7 @@ class ViTEncoder(nn.Module):
 
         if pooling_type is None:
             pooling_out_channels = cnn_activation_out_channels
-            self.pool = nn.Sequential(
-                *[
-                    Rearrange(
-                        f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
-                    ),  # for transformer
-                ]
-            )
+            self.pool = nn.Identity()
 
         elif pooling_type == "max":
             pooling_out_channels = cnn_activation_out_channels
@@ -279,20 +330,30 @@ class ViTEncoder(nn.Module):
                 transformer_embedding_size,
                 transformer_layers,
                 transformer_heads,
-                position_embedding_type=transformer_position_embedding,
+                absolute_position_embedding=transformer_absolute_position_embedding,
+                relative_position_embedding=transformer_relative_position_embedding,
                 source_size=pooling_output_size,
                 mlp_ratio=transformer_mlp_ratio,
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
-                mlp_dropout=transformer_mlp_dropout,
+                ff_linear_module_up=transformer_ff_linear_module_up,
+                ff_linear_module_down=transformer_ff_linear_module_down,
+                msa_scaling=transformer_msa_scaling,
+                ff_dropout=transformer_ff_dropout,
+                ff_inner_dropout=transformer_ff_inner_dropout,
+                ff_outer_dropout=transformer_ff_outer_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
                 causal=False,
                 linear_module=linear_module,
-                bos_tokens=transformer_bos_tokens,
-                return_bos_tokens=transformer_return_bos_tokens,
+                utility_tokens=transformer_utility_tokens,
+                talking_heads=transformer_talking_heads,
+                return_utility_tokens=transformer_return_utility_tokens,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
+                post_norm=transformer_post_norm,
+                checkpoint_ff=transformer_checkpoint_ff,
+                layerscale=transformer_layerscale,
             )
         else:
             self.transformer = nn.Identity()
@@ -304,11 +365,41 @@ class ViTEncoder(nn.Module):
                 transformer_embedding_size,
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
-                dropout=transformer_mlp_dropout,
-                linear_module=linear_module,
+                dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_dropout
+                    if transformer_initial_ff_dropout is not None
+                    else transformer_ff_dropout
+                ),
+                inner_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_inner_dropout
+                    if transformer_initial_ff_inner_dropout is not None
+                    else transformer_ff_inner_dropout
+                ),
+                outer_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_outer_dropout
+                    if transformer_initial_ff_outer_dropout is not None
+                    else transformer_ff_outer_dropout
+                ),
+                linear_module_up=(
+                    # First truthy assigned value
+                    transformer_initial_ff_linear_module_up
+                    or transformer_ff_linear_module_up
+                    or linear_module
+                ),
+                linear_module_down=(
+                    # First truthy assigned value
+                    transformer_initial_ff_linear_module_down
+                    or transformer_ff_linear_module_down
+                    or 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,
+                checkpoint=transformer_checkpoint_ff,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -328,17 +419,24 @@ class ViTEncoder(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         return self.encoder(x)
 
+    def attention_logits(self, x):
+        x = self.encoder[:-1](x)
+        return self.encoder[-1].attention_logits(x)
+
+    def reset_parameters(self):
+        for module in self.encoder:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class ViT(nn.Module):
     """
-    Denoising convolutional transformer
-    Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
-        *''Escaping the Big Data Paradigm with Compact Transformers''*](
-        https://arxiv.org/abs/2104.05704). It's a convolutional neural network
-        leading into a transformer encoder, followed by a sequence pooling layer.
+    ...
     """
 
     def __init__(
@@ -362,23 +460,40 @@ 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_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
-        transformer_position_embedding="relative",  # absolute or relative
+        transformer_post_norm=False,
+        transformer_absolute_position_embedding=False,
+        transformer_relative_position_embedding=True,
         transformer_embedding_size=256,
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=0,
-        transformer_return_bos_tokens=False,
+        transformer_utility_tokens=0,
+        transformer_talking_heads=False,
+        transformer_return_utility_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
-        transformer_mlp_dropout=0.0,
+        transformer_ff_linear_module_up=None,
+        transformer_ff_linear_module_down=None,
+        transformer_msa_scaling="d",
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
-        batch_norm_outputs=True,
-        linear_module=SpectralNormLinear,
+        transformer_checkpoint_ff=True,
+        transformer_layerscale=True,
         head=SequencePoolClassificationHead,
+        batch_norm_logits=True,
+        logit_projection_layer=nn.Linear,
+        linear_module=nn.Linear,
     ):
 
         super().__init__()
@@ -418,33 +533,76 @@ 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_initial_ff_dropout=transformer_initial_ff_dropout,
+            transformer_initial_ff_inner_dropout=transformer_initial_ff_inner_dropout,
+            transformer_initial_ff_outer_dropout=transformer_initial_ff_outer_dropout,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
-            transformer_position_embedding=transformer_position_embedding,
+            transformer_post_norm=transformer_post_norm,
+            transformer_absolute_position_embedding=transformer_absolute_position_embedding,
+            transformer_relative_position_embedding=transformer_relative_position_embedding,
             transformer_embedding_size=transformer_embedding_size,
             transformer_layers=transformer_layers,
             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_utility_tokens=transformer_utility_tokens,
+            transformer_talking_heads=transformer_talking_heads,
+            transformer_return_utility_tokens=transformer_return_utility_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,
-            transformer_mlp_dropout=transformer_mlp_dropout,
+            transformer_ff_linear_module_up=transformer_ff_linear_module_up,
+            transformer_ff_linear_module_down=transformer_ff_linear_module_down,
+            transformer_msa_scaling=transformer_msa_scaling,
+            transformer_ff_dropout=transformer_ff_dropout,
+            transformer_ff_inner_dropout=transformer_ff_inner_dropout,
+            transformer_ff_outer_dropout=transformer_ff_outer_dropout,
             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,
         )
 
         self.pool = head(
             transformer_embedding_size,
-            linear_module,
             image_classes,
-            batch_norm=batch_norm_outputs,
+            logit_projection_layer=logit_projection_layer,
+            batch_norm_logits=batch_norm_logits,
         )
 
+        self.reset_parameters()
+
     @property
     def sequence_length(self):
         return self.encoder.sequence_length
 
     def forward(self, x):
         return self.pool(self.encoder(x))
+
+    def attention_logits(self, x):
+        return self.encoder.attention_logits(x)
+
+    def pool_attention(self, x):
+        if hasattr(self.pool.summarize, "attention"):
+            return self.pool.summarize.attention(self.encoder(x))
+        else:
+            raise NotImplementedError(
+                "`pool_attention` is currently only implemented where"
+                " head class is SequencePoolClassificationHead"
+            )
+
+    def head_to_utility_token_attention_logits(self, x):
+        all_attention = self.attention_logits(x)
+        batch_averages = torch.mean(all_attention, dim=0, keepdim=False)
+        sequence_averages = torch.mean(batch_averages, dim=-1, keepdim=False)
+        n_utility_tokens = self.encoder.encoder[-1]._utility_tokens
+        return sequence_averages[
+            :, :, :n_utility_tokens
+        ]  # (layer, head, utility_tokens)
+
+    def reset_parameters(self):
+        self.encoder.reset_parameters()
+        self.pool.reset_parameters()

{broccoli_ml-0.29.1.dist-info → broccoli_ml-10.0.1.dist-info}/METADATA

@@ -1,17 +1,19 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 0.29.1
+Version: 10.0.1
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey
-Requires-Python: >=3.11
+Requires-Python: >=3.8
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: einops (>=0.8.1,<0.9.0)
-Requires-Dist: numpy (>=2.0.2,<2.1.0)
 Description-Content-Type: text/markdown
 
 # broccoli

broccoli_ml-10.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
+broccoli/activation.py,sha256=nrpTOrpg9k23_E4AJWy7VlXXAJCtCJCOR-TonEWJr04,3218
+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=lnfiv7UIYbABiClIluy6CefGxaiYMrvBcj2Ul0uU6xE,27693
+broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
+broccoli/vit.py,sha256=EGbQb-atuzG3JAx7kdTaJEbWvQR-4XgyYvwjKkN5C38,22612
+broccoli_ml-10.0.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-10.0.1.dist-info/METADATA,sha256=65GKe2Jor5jgUZ8zxROntJ_t0XwAlaukrvpT7nxS0lQ,1369
+broccoli_ml-10.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-10.0.1.dist-info/RECORD,,

broccoli/eigenpatches.py

@@ -1,49 +0,0 @@
-"""
-Jordan (2024) was able to train a CNN to 94% accuracy on CIFAR-10 in 3.29 seconds
-on a single A100 GPU by using carefully-tuned hyperparameters and a number of
-techniques to increase learning efficiency. The author notes that applying fixed
-weights to the first layer of the network that approximate a whitening
-transformation on image patches, following tsyam-code, (2023), was "the single
-most impactful feature... [and] more than doubles training speed".
-
-The usefulness of a fixed layer that whitens image patches can be justified
-according to the work of Chowers & Weiss (2022), who find that the first layer
-weights of a convolutional neural network will asymptotically approach a whitening
-transformation regardless of the details of the rest of the network architecture
-or the training data. This effectively functions as a bandpass filter layer,
-reminiscent of the way neurons in the human primary visual cortex work (Kristensen
-& Sandberg, 2021).
-
-The `eigenvectors` function here is adapted from
-    https://github.com/KellerJordan/cifar10-airbench/blob/master/airbench96_faster.py
-    using https://datascienceplus.com/understanding-the-covariance-matrix/
-"""
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-
-
-def eigenvectors(images: torch.Tensor, patch_size: int, eps=5e-4) -> torch.Tensor:
-    """
-    Adapted from
-        github.com/KellerJordan/cifar10-airbench/blob/master/airbench96_faster.py
-        using https://datascienceplus.com/understanding-the-covariance-matrix/
-
-    Args:
-        images: a batch of training images (the bigger and more representative the better!)
-        patch_size: the size of the eigenvectors we want to create (i.e. the patch/kernel
-            size of the model we will initialise with the eigenvectors)
-        eps: a small number to avoid division by zero
-    """
-    with torch.no_grad():
-        unfolder = nn.Unfold(kernel_size=patch_size, stride=1)
-        patches = unfolder(images)  # (N, patch_elements, patches_per_image)
-        patches = rearrange(patches, "N elements patches -> (N patches) elements")
-        n = patches.size(0)
-        centred = patches - patches.mean(dim=1, keepdim=True)
-        covariance_matrix = (
-            centred.T @ centred
-        ) / n  # https://datascienceplus.com/understanding-the-covariance-matrix/
-        _, eigenvectors = torch.linalg.eigh(covariance_matrix)
-        return eigenvectors

broccoli_ml-0.29.1.dist-info/RECORD

@@ -1,17 +0,0 @@
-broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
-broccoli/activation.py,sha256=-Jf30C6iGqWCorC9HEGn2oduWwjeaCAxGLUUYIy1zX8,3438
-broccoli/assets/2025_resnet_imagenet_1k_pretrained_state_dict.pkl,sha256=RZpPupWxFaVfgZrK-gBgfW1hj78oMEGhVWTbjRB3qMo,46835797
-broccoli/assets/cifar100_eigenvectors_size_2.pt,sha256=DjXDOXMeuMpIqNuGhX9z-OWYVqZwIMScSXZApRr9JjU,2501
-broccoli/assets/cifar100_eigenvectors_size_3.pt,sha256=gL6k0xtXYiYP6ZSvEiMBdJ7kIkT0AngTpDJHFQqwgxA,7173
-broccoli/cnn.py,sha256=jeRyKIAMWu1E3iyI14MGgSZuZivPMh12iqkqW9ilNjo,17785
-broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
-broccoli/linear.py,sha256=8Y9vD85ZEgNZsIQgO3uRQ3lOQR-JjwvabY8liCrfNCk,4831
-broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
-broccoli/tensor.py,sha256=zhSOo9W24FEgN7U35wy3ZIJHnw3u4cepJO5heCw6vwU,4590
-broccoli/transformer.py,sha256=jQGpj_e5WAEU_zEPjCU0OyD_08O3HwsMBg3pbrCzp4E,16924
-broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=m4Wa8B8L25xSODh91ViVyLmwLOBZayp5S7S9f8pIvZo,16109
-broccoli_ml-0.29.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.29.1.dist-info/METADATA,sha256=AojdLjmBwqW9of7D6RSjCdWbDRJO2GM3NbXd1TYeOJY,1257
-broccoli_ml-0.29.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.29.1.dist-info/RECORD,,