broccoli-ml 0.4.3__py3-none-any.whl → 0.5.0__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,

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''*](
@@ -80,6 +80,7 @@ class DCTEncoder(nn.Module):
         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,
@@ -160,7 +161,7 @@ class DCTEncoder(nn.Module):
         if pooling_type in ["maxpool", None]:
             cnn_out_channels = transformer_embedding_size
         elif pooling_type == "concat":
-            cnn_out_channels = min(
+            cnn_out_channels = max(
                 math.floor(transformer_embedding_size / pooling_kernel_voxels),
                 minimum_cnn_out_channels,
             )
@@ -248,14 +249,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 +266,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 +282,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 +303,7 @@ class DCTEncoder(nn.Module):
         return self.encoder(x)
 
 
-class DCT(nn.Module):
+class ViT(nn.Module):
     """
     Denoising convolutional transformer
     Based on the Compact Convolutional Transformer (CCT) of [Hasani et al. (2021)
@@ -321,6 +329,7 @@ class DCT(nn.Module):
         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,12 +338,12 @@ 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,
-        linear_module=nn.Linear,
         initial_batch_norm=True,
+        linear_module=nn.Linear,
         image_classes=100,
     ):
 
@@ -356,7 +365,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 +381,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 +390,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.3.dist-info → broccoli_ml-0.5.0.dist-info}/METADATA

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

{broccoli_ml-0.4.3.dist-info → broccoli_ml-0.5.0.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=-b7XG51c0ZgDoXQRyEmrZH6IvBCH1etK0NTApwqNhpU,16302
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-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,,
+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,,