broccoli-ml 3.2.0__tar.gz → 3.3.1__tar.gz

{broccoli_ml-3.2.0 → broccoli_ml-3.3.1}/PKG-INFO

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

{broccoli_ml-3.2.0 → broccoli_ml-3.3.1}/broccoli/transformer.py

@@ -259,6 +259,13 @@ class FeedforwardBlock(nn.Module):
         self.residual_path = residual_path
         self.post_norm = post_norm
 
+        if self.residual_path and (output_features < input_features):
+            raise ValueError(
+                "If the number of output features will be less than "
+                "the number of input features, then `residual_path` "
+                "should be set to False."
+            )
+
         if self.post_norm:
             self.layernorm = nn.LayerNorm(output_features)
 

broccoli_ml-3.3.1/broccoli/utils.py

@@ -0,0 +1,15 @@
+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)

{broccoli_ml-3.2.0 → broccoli_ml-3.3.1}/broccoli/vit.py

@@ -4,25 +4,12 @@ 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 .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)
 
 
 class GetCLSToken(nn.Module):
@@ -58,7 +45,6 @@ class ClassificationHead(nn.Module):
         self,
         d_model,
         n_classes,
-        linear_module=nn.Linear,
         logit_projection_layer=nn.Linear,
         batch_norm_logits=True,
     ):
@@ -66,10 +52,11 @@ class ClassificationHead(nn.Module):
         self.d_model = d_model
         self.summarize = GetCLSToken()
 
-        if logit_projection_layer is not None:
-            self.projection = logit_projection_layer(d_model, n_classes)
-        else:
+        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)
@@ -99,19 +86,17 @@ class SequencePoolClassificationHead(ClassificationHead):
         self,
         d_model,
         n_classes,
-        linear_module=nn.Linear,
         logit_projection_layer=nn.Linear,
         batch_norm_logits=True,
     ):
         super().__init__(
             d_model,
             n_classes,
-            linear_module=linear_module,
             logit_projection_layer=logit_projection_layer,
             batch_norm_logits=batch_norm_logits,
         )
 
-        self.summarize = SequencePool(d_model, linear_module)
+        self.summarize = SequencePool(d_model, logit_projection_layer)
         # Rebuild the classification process with the correct summary module:
         self.classification_process = nn.Sequential(
             *[
@@ -513,7 +498,6 @@ class ViT(nn.Module):
         self.pool = head(
             transformer_embedding_size,
             image_classes,
-            linear_module=linear_module,
             logit_projection_layer=logit_projection_layer,
             batch_norm_logits=batch_norm_logits,
         )

{broccoli_ml-3.2.0 → broccoli_ml-3.3.1}/pyproject.toml

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