broccoli-ml 4.0.1__tar.gz → 5.0.0__tar.gz

{broccoli_ml-4.0.1 → broccoli_ml-5.0.0}/PKG-INFO

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

{broccoli_ml-4.0.1 → broccoli_ml-5.0.0}/broccoli/transformer.py

@@ -285,9 +285,11 @@ class FeedforwardBlock(nn.Module):
         normformer=False,
         post_norm=True,
         residual_path=True,
+        checkpoint=True,
     ):
         super().__init__()
 
+        self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
 
@@ -326,12 +328,18 @@ class FeedforwardBlock(nn.Module):
         )
 
     def forward(self, x):
+
+        if self.checkpoint:
+            processed = checkpoint(self.process, x, use_reentrant=False)
+        else:
+            processed = self.process(x)
+
         if self.residual_path and self.post_norm:
-            return self.layernorm(x + self.process(x))
+            return self.layernorm(x + processed)
         elif self.residual_path:
-            return x + self.process(x)
+            return x + processed
         else:
-            return self.process(x)
+            return processed
 
 
 class TransformerBlock(nn.Module):
@@ -365,6 +373,7 @@ class TransformerBlock(nn.Module):
         pre_norm=True,
         post_norm=False,
         normformer=False,
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -433,6 +442,7 @@ class TransformerBlock(nn.Module):
             normformer=normformer,
             post_norm=False,  # Handled outside the block
             residual_path=False,  # Handled outside the block
+            checkpoint=checkpoint_ff,
         )
 
     @property
@@ -445,17 +455,17 @@ class TransformerBlock(nn.Module):
             x = self.layer_norm_1(x)
             x = x + self.drop_path(self.attn(x, x, x))
             x = self.layer_norm_2(x)
-            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
             if self.post_norm:  # i.e. in addition! Pre and post.
                 x = self.layer_norm_3(x)
         elif self.post_norm:  # i.e. only, not prenorm, just post
             x = x + self.drop_path(self.attn(x, x, x))
             x = self.layer_norm_1(x)
-            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
             x = self.layer_norm_2(x)
         else:  # Not pre or post norm. Stand well back.
             x = x + self.drop_path(self.attn(x, x, x))
-            x = x + self.drop_path(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
 
         return x
 
@@ -491,6 +501,7 @@ class TransformerEncoder(nn.Module):
         post_norm=False,
         normformer=False,
         msa_scaling="d",
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -567,6 +578,7 @@ class TransformerEncoder(nn.Module):
                     pre_norm=pre_norm,
                     post_norm=post_norm,
                     normformer=normformer,
+                    checkpoint_ff=checkpoint_ff,
                 )
                 for i in range(n_layers)
             ]

{broccoli_ml-4.0.1 → broccoli_ml-5.0.0}/broccoli/vit.py

@@ -160,6 +160,7 @@ class ViTEncoder(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -321,6 +322,7 @@ class ViTEncoder(nn.Module):
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
+                checkpoint_ff=transformer_checkpoint_ff,
             )
         else:
             self.transformer = nn.Identity()
@@ -354,6 +356,7 @@ class ViTEncoder(nn.Module):
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 residual_path=transformer_initial_ff_residual_path,
+                checkpoint=transformer_checkpoint_ff,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -426,6 +429,7 @@ class ViT(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         head=SequencePoolClassificationHead,
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
@@ -492,6 +496,7 @@ class ViT(nn.Module):
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
+            transformer_checkpoint_ff=transformer_checkpoint_ff,
             linear_module=linear_module,
         )
 

{broccoli_ml-4.0.1 → broccoli_ml-5.0.0}/pyproject.toml

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