broccoli-ml 0.5.5__py3-none-any.whl → 0.7.0__py3-none-any.whl

broccoli/transformer.py

@@ -223,7 +223,7 @@ class MHAttention(nn.Module):
 
 class FeedforwardLayer(nn.Module):
     """
-    A denoising autoencoder, of the type used in transformer blocks.
+    ...
     """
 
     def __init__(
@@ -247,6 +247,7 @@ class FeedforwardLayer(nn.Module):
 
         self.process = nn.Sequential(
             *[
+                nn.LayerNorm(input_features),
                 linear_module(
                     input_features,
                     (
@@ -256,8 +257,8 @@ class FeedforwardLayer(nn.Module):
                     ),
                 ),
                 self.activation,
-                self.dropout,
                 linear_module(ratio * input_features, output_features),
+                self.dropout,
             ]
         )
 
@@ -323,25 +324,14 @@ class TransformerBlock(nn.Module):
         )
 
         # Submodules for the feedforward process
-        self.ff_process = nn.Sequential(
-            OrderedDict(
-                [
-                    ("layer_norm", nn.LayerNorm(d_model)),
-                    (
-                        "denoising_autoencoder",
-                        FeedforwardLayer(
-                            d_model,
-                            mlp_ratio,
-                            d_model,
-                            activation=activation,
-                            activation_kwargs=activation_kwargs,
-                            dropout=0.0,
-                            linear_module=linear_module,
-                        ),
-                    ),
-                    ("dropout", nn.Dropout(mlp_dropout)),
-                ]
-            )
+        self.ff = FeedforwardLayer(
+            d_model,
+            mlp_ratio,
+            d_model,
+            activation=activation,
+            activation_kwargs=activation_kwargs,
+            dropout=mlp_dropout,
+            linear_module=linear_module,
         )
 
     @property
@@ -366,7 +356,7 @@ class TransformerBlock(nn.Module):
         process_x = process_x + self.attn(
             norm_process_x, norm_process_x, norm_process_x
         )
-        process_x = process_x + self.ff_process(process_x)
+        process_x = process_x + self.ff(process_x)
         x = torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
 
         return x

broccoli/vit.py

@@ -269,6 +269,13 @@ class ViTEncoder(nn.Module):
                     Rearrange(  # for transformer
                         f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
                     ),
+                    (
+                        PadTensor(
+                            (0, transformer_embedding_size - pooling_out_channels)
+                        )
+                        if not intermediate_feedforward_layer
+                        else nn.Identity()
+                    ),
                 ]
             )
 

{broccoli_ml-0.5.5.dist-info → broccoli_ml-0.7.0.dist-info}/METADATA

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

{broccoli_ml-0.5.5.dist-info → broccoli_ml-0.7.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=SwvutiYOiPlqLzbO_twye7Hna5DsJukVOzzAx9CTCyU,16417
+broccoli/transformer.py,sha256=RSZpbHs_K4ts5os6lWxcGDI3p0zreRwQNnk6mV8HJnk,15930
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=WBYLPowGrt4MUbFGmJiP35F-3wG7CPhDM4lOPgVJy3Q,15365
-broccoli_ml-0.5.5.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.5.5.dist-info/METADATA,sha256=J6BFEZbQdWYaE9raTxH0x05HlTSOnbx365CfGVs4WhU,1256
-broccoli_ml-0.5.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.5.5.dist-info/RECORD,,
+broccoli/vit.py,sha256=_oL0NRUJakyIke2g8WK5eWaiEh06gAhI67l6Wl7k1oM,15659
+broccoli_ml-0.7.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.7.0.dist-info/METADATA,sha256=1QUwYpIruYYiYcMHSgj5lCf-i-FaiipD_5KSAJZeb2s,1256
+broccoli_ml-0.7.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.7.0.dist-info/RECORD,,