broccoli-ml 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

broccoli/vit.py

@@ -53,12 +53,20 @@ class ClassificationHead(nn.Module):
     A general classification head for a ViT
     """
 
-    def __init__(self, d_model, linear_module, n_classes, batch_norm=True):
+    def __init__(
+        self,
+        d_model,
+        n_classes,
+        linear_module=nn.Linear,
+        logit_projection_layer=nn.Linear,
+        batch_norm_logits=True,
+    ):
         super().__init__()
         self.d_model = d_model
         self.summarize = GetCLSToken()
-        self.projection = linear_module(d_model, n_classes)
-        if batch_norm:
+        self.projection = logit_projection_layer(d_model, n_classes)
+
+        if batch_norm_logits:
             self.batch_norm = nn.BatchNorm1d(n_classes, affine=False)
         else:
             self.batch_norm = nn.Identity()
@@ -83,7 +91,7 @@ class SequencePoolClassificationHead(ClassificationHead):
     """
 
     def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
-        super().__init__(d_model, linear_module, out_dim, batch_norm=batch_norm)
+        super().__init__(d_model, linear_module, out_dim, batch_norm_logits=batch_norm)
         self.summarize = SequencePool(d_model, linear_module)
         # Rebuild the classification process with the correct summary module:
         self.classification_process = nn.Sequential(
@@ -411,9 +419,10 @@ class ViT(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
-        batch_norm_outputs=True,
-        linear_module=SpectralNormLinear,
         head=SequencePoolClassificationHead,
+        batch_norm_logits=True,
+        logit_projection_layer=nn.Linear,
+        linear_module=nn.Linear,
     ):
 
         super().__init__()
@@ -480,9 +489,10 @@ class ViT(nn.Module):
 
         self.pool = head(
             transformer_embedding_size,
-            linear_module,
             image_classes,
-            batch_norm=batch_norm_outputs,
+            linear_module=linear_module,
+            logit_projection_layer=logit_projection_layer,
+            batch_norm=batch_norm_logits,
         )
 
     @property

{broccoli_ml-1.1.0.dist-info → broccoli_ml-1.2.0.dist-info}/METADATA

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

{broccoli_ml-1.1.0.dist-info → broccoli_ml-1.2.0.dist-info}/RECORD

@@ -10,8 +10,8 @@ broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
 broccoli/transformer.py,sha256=KhwrV5Nz0utXJQKFFf6MIifmymgOaylc_2_hs9kQt7g,18455
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=7vGg0yHC9tYUo7Rl3rSRCPqdWCUsWEpr06mlo63hWrk,18237
-broccoli_ml-1.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-1.1.0.dist-info/METADATA,sha256=JuVjXWpDhMhRbKrGeMsqHna77AC7z4Vps1fcAQc_pj0,1256
-broccoli_ml-1.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-1.1.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=9dfOPO7y-o96OYpy_ancH4-Y2_Y62fqgzJvI3-86W_o,18471
+broccoli_ml-1.2.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-1.2.0.dist-info/METADATA,sha256=am67vHH2HDQdq3G336ng56xX6-ys1LyZzLRIWaf8dE8,1256
+broccoli_ml-1.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-1.2.0.dist-info/RECORD,,