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

broccoli/transformer.py

@@ -221,7 +221,7 @@ class MHAttention(nn.Module):
         return self.out_proj(output_without_heads)
 
 
-class DenoisingAutoEncoder(nn.Module):
+class FeedforwardLayer(nn.Module):
     """
     A denoising autoencoder, of the type used in transformer blocks.
     """
@@ -329,7 +329,7 @@ class TransformerBlock(nn.Module):
                     ("layer_norm", nn.LayerNorm(d_model)),
                     (
                         "denoising_autoencoder",
-                        DenoisingAutoEncoder(
+                        FeedforwardLayer(
                             d_model,
                             mlp_ratio,
                             d_model,
@@ -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

@@ -1,7 +1,7 @@
 import math
 from typing import Optional
 
-from .transformer import TransformerEncoder, DenoisingAutoEncoder
+from .transformer import TransformerEncoder, FeedforwardLayer
 from .cnn import SpaceToDepth, calculate_output_spatial_size, spatial_tuple
 from .activation import ReLU, SquaredReLU, GELU, SwiGLU
 from einops import einsum
@@ -53,7 +53,7 @@ class SequencePool(nn.Module):
         return self.norm(projection) if self.batch_norm else projection
 
 
-class DCTEncoder(nn.Module):
+class ViTEncoder(nn.Module):
     """
     Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
         *''Escaping the Big Data Paradigm with Compact Transformers''*](
@@ -76,10 +76,11 @@ class DCTEncoder(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,
+        intermediate_feedforward_layer=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -88,8 +89,8 @@ class DCTEncoder(nn.Module):
         transformer_bos_tokens=0,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
-        mlp_dropout=0.0,
-        msa_dropout=0.1,
+        transformer_mlp_dropout=0.0,
+        transformer_msa_dropout=0.1,
         stochastic_depth=0.1,
         linear_module=nn.Linear,
         initial_batch_norm=True,
@@ -113,16 +114,19 @@ class DCTEncoder(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"
@@ -157,7 +161,7 @@ class DCTEncoder(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(
@@ -165,7 +169,9 @@ class DCTEncoder(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
 
@@ -210,7 +216,7 @@ class DCTEncoder(nn.Module):
                 ]
             )
 
-        elif pooling_type == "maxpool":
+        elif pooling_type == "max":
             self.pool = nn.Sequential(
                 *[
                     maxpoolxd(
@@ -224,6 +230,20 @@ class DCTEncoder(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:
@@ -248,14 +268,8 @@ class DCTEncoder(nn.Module):
                     Rearrange(  # for transformer
                         f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
                     ),
-                    DenoisingAutoEncoder(
-                        concatpool_out_channels,
-                        transformer_mlp_ratio,
-                        transformer_embedding_size,
-                        activation=transformer_activation,
-                        activation_kwargs=transformer_activation_kwargs,
-                        dropout=0.0,
-                        linear_module=linear_module,
+                    PadTensor(
+                        (0, transformer_embedding_size - concatpool_out_channels)
                     ),
                 ]
             )
@@ -271,8 +285,8 @@ class DCTEncoder(nn.Module):
                 mlp_ratio=transformer_mlp_ratio,
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
-                mlp_dropout=mlp_dropout,
-                msa_dropout=msa_dropout,
+                mlp_dropout=transformer_mlp_dropout,
+                msa_dropout=transformer_msa_dropout,
                 stochastic_depth=stochastic_depth,
                 causal=False,
                 linear_module=linear_module,
@@ -287,6 +301,19 @@ class DCTEncoder(nn.Module):
                 self.cnn,
                 self.activate_and_dropout,
                 self.pool,
+                (
+                    FeedforwardLayer(
+                        transformer_embedding_size,
+                        transformer_mlp_ratio,
+                        transformer_embedding_size,
+                        activation=transformer_activation,
+                        activation_kwargs=transformer_activation_kwargs,
+                        dropout=transformer_mlp_dropout,
+                        linear_module=linear_module,
+                    )
+                    if intermediate_feedforward_layer
+                    else nn.Identity()
+                ),
                 self.transformer,
             ]
         )
@@ -295,7 +322,7 @@ class DCTEncoder(nn.Module):
         return self.encoder(x)
 
 
-class DCT(nn.Module):
+class CCT(nn.Module):
     """
     Denoising convolutional transformer
     Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
@@ -317,10 +344,11 @@ class DCT(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,
+        intermediate_feedforward_layer=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -329,8 +357,8 @@ class DCT(nn.Module):
         transformer_bos_tokens=0,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
-        mlp_dropout=0.0,
-        msa_dropout=0.1,
+        transformer_mlp_dropout=0.0,
+        transformer_msa_dropout=0.1,
         stochastic_depth=0.1,
         batch_norm_outputs=True,
         initial_batch_norm=True,
@@ -356,7 +384,7 @@ class DCT(nn.Module):
                 "SwiGLU": SwiGLU,
             }[transformer_activation]
 
-        self.encoder = DCTEncoder(
+        self.encoder = ViTEncoder(
             input_size=input_size,
             cnn_in_channels=cnn_in_channels,
             minimum_cnn_out_channels=minimum_cnn_out_channels,
@@ -372,6 +400,7 @@ class DCT(nn.Module):
             pooling_kernel_size=pooling_kernel_size,
             pooling_kernel_stride=pooling_kernel_stride,
             pooling_padding=pooling_padding,
+            intermediate_feedforward_layer=intermediate_feedforward_layer,
             transformer_position_embedding=transformer_position_embedding,
             transformer_embedding_size=transformer_embedding_size,
             transformer_layers=transformer_layers,
@@ -380,8 +409,8 @@ class DCT(nn.Module):
             transformer_bos_tokens=transformer_bos_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,
-            mlp_dropout=mlp_dropout,
-            msa_dropout=msa_dropout,
+            mlp_dropout=transformer_mlp_dropout,
+            msa_dropout=transformer_msa_dropout,
             stochastic_depth=stochastic_depth,
             linear_module=linear_module,
             initial_batch_norm=initial_batch_norm,

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

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

{broccoli_ml-0.4.4.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=23R58t3TLZMb9ulhCtQ3gXu0mPlfyPvLM8TaGOpaz58,16310
+broccoli/transformer.py,sha256=SwvutiYOiPlqLzbO_twye7Hna5DsJukVOzzAx9CTCyU,16417
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=FQbdwNpvdqHreTdh_rK98sqhIrVVtn4L77EmNSvxXK0,14086
-broccoli_ml-0.4.4.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.4.4.dist-info/METADATA,sha256=rsyvyeUbOLY_QLuC-klSHxaLVWZjfffKf2ugDuJXnR8,1256
-broccoli_ml-0.4.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.4.4.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,,