broccoli-ml 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

broccoli/transformer.py

@@ -255,20 +255,24 @@ class FeedforwardBlock(nn.Module):
 
         if regularise_values:
             self.memory_type = SpectralNormLinear
-            self.bias_memories = False
         else:
             self.memory_type = nn.Linear
-            self.bias_memories = True
 
         self.process = nn.Sequential(
             *[
-                nn.LayerNorm(input_features),
+                (
+                    nn.LayerNorm(input_features)
+                    if not self.regularise_values
+                    else nn.Identity()
+                ),
                 linear_module(input_features, self.max_features),
                 self.activation,
-                nn.LayerNorm(ratio * output_features),
-                self.memory_type(
-                    ratio * output_features, output_features, bias=self.bias_memories
+                (
+                    nn.LayerNorm(input_features)
+                    if not self.regularise_values
+                    else nn.Identity()
                 ),
+                self.memory_type(ratio * output_features, output_features),
                 self.dropout,
             ]
         )

broccoli/vit.py

@@ -30,7 +30,7 @@ class GetCLSToken(nn.Module):
 
 
 class SequencePool(nn.Module):
-    def __init__(self, d_model, linear_module):
+    def __init__(self, d_model, linear_module=nn.Linear):
         super().__init__()
         self.attention = nn.Sequential(
             *[
@@ -54,13 +54,6 @@ class ClassificationHead(nn.Module):
         super().__init__()
         self.d_model = d_model
         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, n_classes)
         if batch_norm:
             self.batch_norm = nn.BatchNorm1d(n_classes, affine=False)
@@ -87,7 +80,7 @@ class SequencePoolClassificationHead(ClassificationHead):
     """
 
     def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
-        super().__init__(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)
 
 

{broccoli_ml-0.19.0.dist-info → broccoli_ml-0.21.0.dist-info}/METADATA

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

{broccoli_ml-0.19.0.dist-info → broccoli_ml-0.21.0.dist-info}/RECORD

@@ -8,10 +8,10 @@ broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=g8YrxNl6g_WcHrWVmbaBHJU5hv6daFS0r4TxAoPJ9UE,3012
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=MUvXtwD2f1sPTBym4FB0x_ZfsJUBNLgULUlN8btV8GI,1943
-broccoli/transformer.py,sha256=2hoRg9N0g2HTfrUU6cTXl4msHO0orWiHS9xJTUMHej4,16234
+broccoli/transformer.py,sha256=2O9EALlZPQXYCtOGvkCgoWv0-K8rNtsp-MTon1h_4n8,16343
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=--eEAxBzCCVGCIvaZWHsonbVXGe04wDEz0Q3V4YZXoI,16599
-broccoli_ml-0.19.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.19.0.dist-info/METADATA,sha256=aSUin_G9VrKRhl8gBqGirK1yWJyMo6krnWzA0KWF1tQ,1257
-broccoli_ml-0.19.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.19.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=vQtOcC0Dd8y6PTWx0xCbnE4ymYkL_HfYrerqaJ0hs1k,16404
+broccoli_ml-0.21.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.21.0.dist-info/METADATA,sha256=jtoA3JoVE-R_VYlofiCzNxnWtzvV1Tu1HjR_Bm9T8lA,1257
+broccoli_ml-0.21.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.21.0.dist-info/RECORD,,