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

broccoli/vit.py

@@ -3,9 +3,21 @@ from typing import Optional
 
 from .transformer import TransformerEncoder
 from .cnn import ConvLayer, ConcatPool
+from .activation import ReLU, SquaredReLU, GELU, SwiGLU
 from einops import einsum
 from einops.layers.torch import Rearrange
 import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PadTensor(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.kwargs = kwargs
+
+    def forward(self, x):
+        return F.pad(x, *self.args, **self.kwargs)
 
 
 class SequencePool(nn.Module):
@@ -109,7 +121,7 @@ class CCTEncoder(nn.Module):
             conv_out_channels = transformer_embedding_size
         elif conv_pooling_type == "concat":
             conv_out_channels = int(
-                round(transformer_embedding_size / (conv_pooling_kernel_size**2))
+                math.floor(transformer_embedding_size / (conv_pooling_kernel_size**2))
             )
 
         # This if block rhymes:
@@ -143,11 +155,15 @@ class CCTEncoder(nn.Module):
             )
 
         elif conv_pooling_type == "concat":
-            concatpool_activation_output_size = (
+            concatpool_activation_output_channels = (
                 conv_pooling_kernel_size**2 * conv_out_channels
             )
             if cnn_activation.__name__.endswith("GLU"):
-                concatpool_activation_output_size /= 2
+                concatpool_activation_output_channels /= 2
+
+            concatpool_padding = (
+                transformer_embedding_size - concatpool_activation_output_channels
+            )
 
             self.pool = nn.Sequential(
                 *[
@@ -161,10 +177,7 @@ class CCTEncoder(nn.Module):
                     ),
                     self.cnn_activation,
                     Rearrange("N H W C -> N (H W) C"),
-                    nn.Linear(
-                        concatpool_activation_output_size, transformer_embedding_size
-                    ),
-                    self.cnn_activation,
+                    PadTensor((0, concatpool_padding)),
                 ]
             )
 
@@ -249,6 +262,22 @@ class CCT(nn.Module):
 
         super().__init__()
 
+        if isinstance(cnn_activation, str):
+            cnn_activation = {
+                "ReLU": ReLU,
+                "SquaredReLU": SquaredReLU,
+                "GELU": GELU,
+                "SwiGLU": SwiGLU,
+            }[cnn_activation]
+
+        if isinstance(transformer_activation, str):
+            transformer_activation = {
+                "ReLU": ReLU,
+                "SquaredReLU": SquaredReLU,
+                "GELU": GELU,
+                "SwiGLU": SwiGLU,
+            }[transformer_activation]
+
         self.encoder = CCTEncoder(
             image_size=image_size,
             conv_kernel_size=conv_kernel_size,

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

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

{broccoli_ml-0.1.36.dist-info → broccoli_ml-0.1.38.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=WEIJoemSDbYwZXNL3ArW4WhLIaWz1h7nTK3_LHD5Oyk,10841
-broccoli_ml-0.1.36.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.1.36.dist-info/METADATA,sha256=FXohxMQF2qtcUv9B_JeRXgihoDYVUCqCK5ns-PeSx8o,1257
-broccoli_ml-0.1.36.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.1.36.dist-info/RECORD,,
+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,,