broccoli-ml 0.22.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

broccoli/transformer.py

@@ -236,7 +236,7 @@ class FeedforwardBlock(nn.Module):
         activation_kwargs=None,
         dropout=0.0,
         linear_module=nn.Linear,
-        regularise_values=True,
+        sigma_reparam=False,
     ):
         super().__init__()
 
@@ -253,21 +253,17 @@ class FeedforwardBlock(nn.Module):
             else ratio * output_features
         )
 
-        if regularise_values:
+        if sigma_reparam:
             self.memory_type = SpectralNormLinear
         else:
-            self.memory_type = nn.Linear
+            self.memory_type = linear_module
 
         self.process = nn.Sequential(
             *[
                 nn.LayerNorm(input_features),
                 linear_module(input_features, self.max_features),
                 self.activation,
-                # (
-                #     nn.LayerNorm(input_features)
-                #     if not regularise_values
-                #     else nn.Identity()
-                # ),
+                nn.LayerNorm(ratio * output_features),
                 self.memory_type(ratio * output_features, output_features),
                 self.dropout,
             ]

broccoli/vit.py

@@ -18,7 +18,10 @@ class PadTensor(nn.Module):
         self.kwargs = kwargs
 
     def forward(self, x):
-        return F.pad(x, *self.args, **self.kwargs)
+        if sum(self.args[0]) == 0:
+            return x
+        else:
+            return F.pad(x, *self.args, **self.kwargs)
 
 
 class GetCLSToken(nn.Module):
@@ -97,9 +100,9 @@ class ViTEncoder(nn.Module):
     def __init__(
         self,
         input_size=(32, 32),
+        in_channels=3,
         initial_batch_norm=True,
         cnn=True,
-        cnn_in_channels=3,
         cnn_out_channels=16,
         cnn_kernel_size=3,
         cnn_kernel_stride=1,
@@ -113,7 +116,7 @@ class ViTEncoder(nn.Module):
         pooling_kernel_size=3,
         pooling_kernel_stride=2,
         pooling_padding=1,
-        intermediate_feedforward_layer=True,
+        transformer_feedforward_first=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -180,7 +183,7 @@ class ViTEncoder(nn.Module):
                 dilation=cnn_kernel_dilation,
             )
             self.cnn = convxd(
-                cnn_in_channels,
+                in_channels,
                 cnn_out_channels,
                 cnn_kernel_size,
                 stride=cnn_kernel_stride,
@@ -206,8 +209,8 @@ class ViTEncoder(nn.Module):
             self.cnn = nn.Identity()
             self.activate_and_dropout = nn.Identity()
             cnn_output_size = input_size
-            cnn_out_channels = cnn_in_channels
-            cnn_activation_out_channels = cnn_in_channels
+            cnn_out_channels = in_channels
+            cnn_activation_out_channels = in_channels
 
         pooling_kernel_voxels = math.prod(
             spatial_tuple(pooling_kernel_size, self.spatial_dimensions)
@@ -262,6 +265,10 @@ class ViTEncoder(nn.Module):
                 "Pooling type must be max, average, concat or None"
             )
 
+        self.pooling_channels_padding = PadTensor(
+            (0, max(0, transformer_embedding_size - pooling_out_channels))
+        )
+
         self.sequence_length = math.prod(pooling_output_size)  # One token per voxel
 
         if transformer_layers > 0:
@@ -286,38 +293,23 @@ class ViTEncoder(nn.Module):
         else:
             self.transformer = nn.Identity()
 
-        if intermediate_feedforward_layer:
-            self.pooling_channels_padding = nn.Identity()
-            self.intermediate_feedforward_layer = FeedforwardBlock(
-                pooling_out_channels,
+        if transformer_feedforward_first:
+            self.initial_ff = FeedforwardBlock(
+                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,
-            )
-        elif pooling_out_channels == transformer_embedding_size:
-            self.intermediate_feedforward_layer = nn.Identity()
-            self.pooling_channels_padding = nn.Identity()
-        elif pooling_out_channels < transformer_embedding_size:
-            self.intermediate_feedforward_layer = nn.Identity()
-            self.pooling_channels_padding = PadTensor(
-                (0, transformer_embedding_size - pooling_out_channels)
+                sigma_reparam=True,
             )
         else:
-            raise NotImplementedError(
-                "In a situation where the choice/parameters of the pooling and the"
-                + " `cnn_out_channels` (or the number of `input_channels` if"
-                + " `cnn`=False) means that the pooling will result"
-                + " in more channels per pixel/voxel than the size of the"
-                + " intended transformer embedding,"
-                + " `intermediate_feedforward_layer` must be set to True"
-            )
+            self.initial_ff = nn.Identity()
 
         self.encoder = nn.Sequential(
             *[
-                batchnormxd(cnn_in_channels) if initial_batch_norm else nn.Identity(),
+                batchnormxd(in_channels) if initial_batch_norm else nn.Identity(),
                 self.cnn,
                 self.activate_and_dropout,
                 self.pool,
@@ -325,7 +317,7 @@ class ViTEncoder(nn.Module):
                     f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
                 ),
                 self.pooling_channels_padding,
-                self.intermediate_feedforward_layer,
+                self.initial_ff,
                 self.transformer,
             ]
         )
@@ -347,9 +339,9 @@ class ViT(nn.Module):
         self,
         input_size=(32, 32),
         image_classes=100,
+        in_channels=3,
         initial_batch_norm=True,
         cnn=True,
-        cnn_in_channels=3,
         cnn_out_channels=16,
         cnn_kernel_size=3,
         cnn_kernel_stride=1,
@@ -363,7 +355,7 @@ class ViT(nn.Module):
         pooling_kernel_size=3,
         pooling_kernel_stride=2,
         pooling_padding=1,
-        intermediate_feedforward_layer=True,
+        transformer_feedforward_first=True,
         transformer_position_embedding="relative",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -402,8 +394,8 @@ class ViT(nn.Module):
         self.encoder = ViTEncoder(
             input_size=input_size,
             initial_batch_norm=initial_batch_norm,
+            in_channels=in_channels,
             cnn=cnn,
-            cnn_in_channels=cnn_in_channels,
             cnn_out_channels=cnn_out_channels,
             cnn_kernel_size=cnn_kernel_size,
             cnn_kernel_stride=cnn_kernel_stride,
@@ -417,7 +409,7 @@ class ViT(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_feedforward_first=transformer_feedforward_first,
             transformer_position_embedding=transformer_position_embedding,
             transformer_embedding_size=transformer_embedding_size,
             transformer_layers=transformer_layers,

{broccoli_ml-0.22.0.dist-info → broccoli_ml-0.23.0.dist-info}/METADATA

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

{broccoli_ml-0.22.0.dist-info → broccoli_ml-0.23.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=LrrHzBRs-WmqajBTdt2BTelfL8Lo-ESyln9dOh9ME8U,16219
+broccoli/transformer.py,sha256=NxOHP-XQRCtoiiTh7WJWNvSjpZzamiqQU966nQh5vhQ,16091
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=vQtOcC0Dd8y6PTWx0xCbnE4ymYkL_HfYrerqaJ0hs1k,16404
-broccoli_ml-0.22.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.22.0.dist-info/METADATA,sha256=5fSF08I5CmTnJr2XMp5qIXY7rWKMBaYl2otbRj3IqLY,1257
-broccoli_ml-0.22.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.22.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=LRofeDyz0wMwClVBeuxnaPmuDpuu8UAVddAOxIOYK8Y,15625
+broccoli_ml-0.23.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.23.0.dist-info/METADATA,sha256=GfhIQiwV5g480TcVdGdzL3AsgJbQmV89Rxc_oJOrwJo,1257
+broccoli_ml-0.23.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.23.0.dist-info/RECORD,,