broccoli-ml 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

broccoli/transformer.py

@@ -395,6 +395,7 @@ class TransformerEncoder(nn.Module):
         causal=False,
         linear_module=nn.Linear,
         bos_tokens=0,
+        return_bos_tokens=False,
     ):
         if position_embedding_type == "relative":
             assert source_size is not None  # TODO: make this a proper exception
@@ -403,6 +404,7 @@ class TransformerEncoder(nn.Module):
         self.seq_len = seq_len
         self.n_heads = n_heads
         self._bos_tokens = bos_tokens
+        self.return_bos_tokens = return_bos_tokens
 
         # Initialise BOS tokens with normal init, like usual Pytorch embeddings
         if self._bos_tokens:
@@ -479,7 +481,7 @@ class TransformerEncoder(nn.Module):
         for block in self.blocks:
             x = block(x)
 
-        if self._bos_tokens:
+        if self._bos_tokens and not self.return_bos_tokens:
             return x[:, self._bos_tokens :, :]
         else:
             return x

broccoli/vit.py

@@ -76,7 +76,7 @@ class ViTEncoder(nn.Module):
         cnn_activation: nn.Module = ReLU,
         cnn_activation_kwargs: Optional[dict] = None,
         cnn_dropout=0.0,
-        pooling_type="concat",  # maxpool or concat
+        pooling_type="concat",  # max, average or concat
         pooling_kernel_size=3,
         pooling_kernel_stride=2,
         pooling_padding=1,
@@ -114,16 +114,19 @@ class ViTEncoder(nn.Module):
 
         if self.spatial_dimensions == 1:
             maxpoolxd = nn.MaxPool1d
+            avgpoolxd = nn.AvgPool1d
             convxd = nn.Conv1d
             batchnormxd = nn.BatchNorm1d
             spatial_dim_names = "D1"
         elif self.spatial_dimensions == 2:
             maxpoolxd = nn.MaxPool2d
+            avgpoolxd = nn.AvgPool2d
             convxd = nn.Conv2d
             batchnormxd = nn.BatchNorm2d
             spatial_dim_names = "D1 D2"
         elif self.spatial_dimensions == 3:
             maxpoolxd = nn.MaxPool3d
+            avgpoolxd = nn.AvgPool3d
             convxd = nn.Conv3d
             batchnormxd = nn.BatchNorm3d
             spatial_dim_names = "D1 D2 D3"
@@ -158,7 +161,7 @@ class ViTEncoder(nn.Module):
             spatial_tuple(pooling_kernel_size, self.spatial_dimensions)
         )
 
-        if pooling_type in ["maxpool", None]:
+        if pooling_type in ["max", "average", None]:
             cnn_out_channels = transformer_embedding_size
         elif pooling_type == "concat":
             cnn_out_channels = max(
@@ -166,7 +169,9 @@ class ViTEncoder(nn.Module):
                 minimum_cnn_out_channels,
             )
         else:
-            raise NotImplementedError("Pooling type must be maxpool, concat or None")
+            raise NotImplementedError(
+                "Pooling type must be max, average, concat or None"
+            )
 
         cnn_activation_out_channels = cnn_out_channels
 
@@ -211,7 +216,7 @@ class ViTEncoder(nn.Module):
                 ]
             )
 
-        elif pooling_type == "maxpool":
+        elif pooling_type == "max":
             self.pool = nn.Sequential(
                 *[
                     maxpoolxd(
@@ -225,6 +230,20 @@ class ViTEncoder(nn.Module):
                 ]
             )
 
+        elif pooling_type == "average":
+            self.pool = nn.Sequential(
+                *[
+                    avgpoolxd(
+                        pooling_kernel_size,
+                        stride=pooling_kernel_stride,
+                        padding=pooling_padding,
+                    ),
+                    Rearrange(
+                        f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
+                    ),  # for transformer
+                ]
+            )
+
         elif pooling_type == "concat":
 
             if transformer_activation_kwargs is not None:
@@ -303,7 +322,7 @@ class ViTEncoder(nn.Module):
         return self.encoder(x)
 
 
-class ViT(nn.Module):
+class CCT(nn.Module):
     """
     Denoising convolutional transformer
     Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
@@ -325,7 +344,7 @@ class ViT(nn.Module):
         cnn_activation: nn.Module = ReLU,
         cnn_activation_kwargs: Optional[dict] = None,
         cnn_dropout=0.0,
-        pooling_type="concat",  # maxpool or concat
+        pooling_type="concat",  # max, average or concat
         pooling_kernel_size=3,
         pooling_kernel_stride=2,
         pooling_padding=1,

{broccoli_ml-0.5.0.dist-info → broccoli_ml-0.5.1.dist-info}/METADATA

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

{broccoli_ml-0.5.0.dist-info → broccoli_ml-0.5.1.dist-info}/RECORD

@@ -8,10 +8,10 @@ broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=0XYCi3ckTEKwAgBOMUSJP2HsnrroOH8eyrhRdpANG2w,1298
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=E2JK5mQwJf75e23-JGcDoT7QxQf89DJReUo2et1LhRY,1716
-broccoli/transformer.py,sha256=-b7XG51c0ZgDoXQRyEmrZH6IvBCH1etK0NTApwqNhpU,16302
+broccoli/transformer.py,sha256=SwvutiYOiPlqLzbO_twye7Hna5DsJukVOzzAx9CTCyU,16417
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=7XJp53TEcGo2jaWL-XtEd_RTAvl2er3DmAcXawm1SkY,14627
-broccoli_ml-0.5.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.5.0.dist-info/METADATA,sha256=eN2lZ82O-87kkq6Kypo2MXzltUUKMvrulUDphy2-lNE,1256
-broccoli_ml-0.5.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.5.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=CSllHj_cI0dQveuLSoHD-Y95nOqVO1F-p-8SmlJffhM,15272
+broccoli_ml-0.5.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.5.1.dist-info/METADATA,sha256=IFo3lnkE6Zti81jUvLuBONLlR2PUUAwj1AKqPCgFUdI,1256
+broccoli_ml-0.5.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.5.1.dist-info/RECORD,,