broccoli-ml 0.1.38__py3-none-any.whl → 0.1.40__py3-none-any.whl

broccoli/vit.py

@@ -67,6 +67,7 @@ class CCTEncoder(nn.Module):
         conv_pooling_kernel_size=3,
         conv_pooling_kernel_stride=2,
         conv_pooling_kernel_padding=1,
+        conv_dropout=0.0,
         transformer_position_embedding="absolute",  # absolute or relative
         transformer_embedding_size=256,
         transformer_layers=7,
@@ -155,15 +156,20 @@ class CCTEncoder(nn.Module):
             )
 
         elif conv_pooling_type == "concat":
-            concatpool_activation_output_channels = (
-                conv_pooling_kernel_size**2 * conv_out_channels
-            )
-            if cnn_activation.__name__.endswith("GLU"):
-                concatpool_activation_output_channels /= 2
 
-            concatpool_padding = (
-                transformer_embedding_size - concatpool_activation_output_channels
-            )
+            if transformer_activation_kwargs is not None:
+                self.concatpool_activation = transformer_activation(
+                    **transformer_activation_kwargs
+                )
+            else:
+                self.concatpool_activation = transformer_activation()
+
+            concatpool_out_channels = conv_pooling_kernel_size**2 * conv_out_channels
+
+            if cnn_activation.__name__.endswith("GLU"):
+                cnn_activation_output_channels = concatpool_out_channels / 2
+            else:
+                cnn_activation_output_channels = concatpool_out_channels
 
             self.pool = nn.Sequential(
                 *[
@@ -176,8 +182,24 @@ class CCTEncoder(nn.Module):
                         "N C H W -> N H W C"
                     ),
                     self.cnn_activation,
-                    Rearrange("N H W C -> N (H W) C"),
-                    PadTensor((0, concatpool_padding)),
+                    nn.Dropout(conv_dropout),
+                    Rearrange(  # rearrange in case we're using XGLU activation
+                        "N H W C -> N C H W"
+                    ),
+                    nn.BatchNorm2d(cnn_activation_output_channels),
+                    Rearrange(  # rearrange in case we're using XGLU activation
+                        "N C H W -> N (H W) C"
+                    ),
+                    nn.Linear(
+                        cnn_activation_output_channels,
+                        (
+                            2 * transformer_embedding_size * transformer_mlp_ratio
+                            if transformer_activation.__name__.endswith("GLU")
+                            else transformer_embedding_size * transformer_mlp_ratio
+                        ),
+                    ),
+                    self.concatpool_activation,
+                    nn.Linear(transformer_embedding_size * transformer_mlp_ratio),
                 ]
             )
 

{broccoli_ml-0.1.38.dist-info → broccoli_ml-0.1.40.dist-info}/METADATA

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

{broccoli_ml-0.1.38.dist-info → broccoli_ml-0.1.40.dist-info}/RECORD

@@ -10,8 +10,8 @@ broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=E2JK5mQwJf75e23-JGcDoT7QxQf89DJReUo2et1LhRY,1716
 broccoli/transformer.py,sha256=gFBIEowGFPSgQhM1RwsRtQlw_WzVJPY-LJyf1MLtPek,16277
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=CSv13ILKw12o1fNznFvgbfw1TR-gDW30h74yjW6HmLc,11692
-broccoli_ml-0.1.38.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.1.38.dist-info/METADATA,sha256=JJpmHolP3y4Yz8oZ5eRSZswLlCBsqxu6Y5VfimaD5c8,1257
-broccoli_ml-0.1.38.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.1.38.dist-info/RECORD,,
+broccoli/vit.py,sha256=rB1hfEwqRDfSFnFgXzZtzIeqFLnDfiyWe6hDJ7OcH8Q,12777
+broccoli_ml-0.1.40.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.1.40.dist-info/METADATA,sha256=eTdLGu8jKvslYTs2_1qQd-GdV5vSOLJUuYDFHD3CEk8,1257
+broccoli_ml-0.1.40.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.1.40.dist-info/RECORD,,