broccoli-ml 0.1.37__tar.gz → 0.1.39__tar.gz

{broccoli_ml-0.1.37 → broccoli_ml-0.1.39}/PKG-INFO

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

{broccoli_ml-0.1.37 → broccoli_ml-0.1.39}/broccoli/vit.py

@@ -7,6 +7,17 @@ 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):
@@ -56,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,
@@ -110,7 +122,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:
@@ -144,11 +156,16 @@ class CCTEncoder(nn.Module):
             )
 
         elif conv_pooling_type == "concat":
-            concatpool_activation_output_size = (
-                conv_pooling_kernel_size**2 * conv_out_channels
+            self.concatpool_activation = transformer_activation(
+                **transformer_activation_kwargs
             )
+
+            concatpool_out_channels = conv_pooling_kernel_size**2 * conv_out_channels
+
             if cnn_activation.__name__.endswith("GLU"):
-                concatpool_activation_output_size /= 2
+                cnn_activation_output_channels = concatpool_out_channels / 2
+            else:
+                cnn_activation_output_channels = concatpool_out_channels
 
             self.pool = nn.Sequential(
                 *[
@@ -161,11 +178,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"),
+                    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(
-                        concatpool_activation_output_size, transformer_embedding_size
+                        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.cnn_activation,
+                    self.concatpool_activation,
+                    nn.Linear(transformer_embedding_size * transformer_mlp_ratio),
                 ]
             )
 

{broccoli_ml-0.1.37 → broccoli_ml-0.1.39}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "0.1.37"
+version = "0.1.39"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}