broccoli-ml 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

broccoli/vit.py

@@ -27,7 +27,7 @@ class SequencePool(nn.Module):
         as a generalisation of average pooling.
     """
 
-    def __init__(self, d_model, linear_module, out_dim):
+    def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
         super().__init__()
         self.d_model = d_model
         self.attention = nn.Sequential(
@@ -38,7 +38,11 @@ class SequencePool(nn.Module):
             ]
         )
         self.projection = nn.Linear(d_model, out_dim)
-        self.norm = nn.BatchNorm1d(out_dim, affine=False)
+        self.batch_norm = batch_norm
+        if batch_norm:
+            self.norm = nn.BatchNorm1d(out_dim, affine=False)
+        else:
+            self.norm = None
 
     def forward(self, x):
         weights = self.attention(x)
@@ -46,7 +50,7 @@ class SequencePool(nn.Module):
             weights, x, "batch seq, batch seq d_model -> batch d_model"
         )
         projection = self.projection(weighted_embedding)
-        return self.norm(projection)
+        return self.norm(projection) if self.batch_norm else projection
 
 
 class DCTEncoder(nn.Module):
@@ -81,14 +85,14 @@ class DCTEncoder(nn.Module):
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=4,
+        transformer_bos_tokens=0,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
         mlp_dropout=0.0,
         msa_dropout=0.1,
         stochastic_depth=0.1,
         linear_module=nn.Linear,
-        batch_norm=True,
+        initial_batch_norm=True,
     ):
         super().__init__()
 
@@ -191,7 +195,7 @@ class DCTEncoder(nn.Module):
                 nn.Dropout(cnn_dropout),
                 (
                     batchnormxd(cnn_activation_out_channels)
-                    if batch_norm
+                    if initial_batch_norm
                     else nn.Identity()
                 ),
             ]
@@ -279,7 +283,7 @@ class DCTEncoder(nn.Module):
 
         self.encoder = nn.Sequential(
             *[
-                batchnormxd(cnn_in_channels) if batch_norm else nn.Identity(),
+                batchnormxd(cnn_in_channels) if initial_batch_norm else nn.Identity(),
                 self.cnn,
                 self.activate_and_dropout,
                 self.pool,
@@ -322,14 +326,15 @@ class DCT(nn.Module):
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=4,
+        transformer_bos_tokens=0,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
         mlp_dropout=0.0,
         msa_dropout=0.1,
         stochastic_depth=0.1,
+        batch_norm_outputs=True,
         linear_module=nn.Linear,
-        batch_norm=True,
+        initial_batch_norm=True,
         image_classes=100,
     ):
 
@@ -379,10 +384,13 @@ class DCT(nn.Module):
             msa_dropout=msa_dropout,
             stochastic_depth=stochastic_depth,
             linear_module=linear_module,
-            batch_norm=batch_norm,
+            initial_batch_norm=initial_batch_norm,
         )
         self.pool = SequencePool(
-            transformer_embedding_size, linear_module, image_classes
+            transformer_embedding_size,
+            linear_module,
+            image_classes,
+            batch_norm=batch_norm_outputs,
         )
 
     @property

{broccoli_ml-0.4.1.dist-info → broccoli_ml-0.4.3.dist-info}/METADATA

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

{broccoli_ml-0.4.1.dist-info → broccoli_ml-0.4.3.dist-info}/RECORD

@@ -10,8 +10,8 @@ broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=E2JK5mQwJf75e23-JGcDoT7QxQf89DJReUo2et1LhRY,1716
 broccoli/transformer.py,sha256=23R58t3TLZMb9ulhCtQ3gXu0mPlfyPvLM8TaGOpaz58,16310
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=d9nKhohlxpFbu3wzhNi53bYBNMUuShPfF6NXUAyDVA0,13778
-broccoli_ml-0.4.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.4.1.dist-info/METADATA,sha256=qwzzby85q__wYdKBEYBgyW5D7q3GMPoRwzeHPY6Mf6s,1256
-broccoli_ml-0.4.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.4.1.dist-info/RECORD,,
+broccoli/vit.py,sha256=s1QwZac3S-QjAFEujf0vDDMacYwr_aWE_1mFvFD18-4,14086
+broccoli_ml-0.4.3.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.4.3.dist-info/METADATA,sha256=KzHrCtuxGAdNbi5ORjQg9KJSdgAK9OtcP6ubw3owZ9Q,1256
+broccoli_ml-0.4.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.4.3.dist-info/RECORD,,