broccoli-ml 0.13.0__py3-none-any.whl → 0.13.1__py3-none-any.whl

broccoli/linear.py

@@ -10,7 +10,8 @@ from .tensor import SigmaReparamTensor
 
 class SpectralNormLinear(nn.Module):
     """
-    ...
+    Inspired by Apple's Spectral Normed Linear Layers
+        (https://github.com/apple/ml-sigma-reparam)
     """
 
     def __init__(self, in_features: int, out_features: int, bias: bool = True):

broccoli/transformer.py

@@ -222,7 +222,7 @@ class MHAttention(nn.Module):
         return self.out_proj(output_without_heads)
 
 
-class FeedforwardLayer(nn.Module):
+class FeedforwardBlock(nn.Module):
     """
     ...
     """
@@ -324,7 +324,7 @@ class TransformerBlock(nn.Module):
         )
 
         # Submodules for the feedforward process
-        self.ff = FeedforwardLayer(
+        self.ff = FeedforwardBlock(
             d_model,
             mlp_ratio,
             d_model,

broccoli/vit.py

@@ -1,9 +1,10 @@
 import math
 from typing import Optional
 
-from .transformer import TransformerEncoder, FeedforwardLayer
+from .transformer import TransformerEncoder, FeedforwardBlock
 from .cnn import SpaceToDepth, calculate_output_spatial_size, spatial_tuple
 from .activation import ReLU, SquaredReLU, GELU, SwiGLU
+from .linear import SpectralNormLinear
 from einops import einsum
 from einops.layers.torch import Rearrange
 import torch.nn as nn
@@ -103,6 +104,7 @@ class ViTEncoder(nn.Module):
     def __init__(
         self,
         input_size=(32, 32),
+        initial_batch_norm=True,
         cnn=True,
         cnn_in_channels=3,
         cnn_out_channels=16,
@@ -132,7 +134,6 @@ class ViTEncoder(nn.Module):
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         linear_module=nn.Linear,
-        initial_batch_norm=True,
     ):
         super().__init__()
 
@@ -294,7 +295,7 @@ class ViTEncoder(nn.Module):
 
         if intermediate_feedforward_layer:
             self.pooling_channels_padding = nn.Identity()
-            self.intermediate_feedforward_layer = FeedforwardLayer(
+            self.intermediate_feedforward_layer = FeedforwardBlock(
                 pooling_out_channels,
                 transformer_mlp_ratio,
                 transformer_embedding_size,
@@ -349,6 +350,8 @@ class ViT(nn.Module):
     def __init__(
         self,
         input_size=(32, 32),
+        image_classes=100,
+        initial_batch_norm=True,
         cnn=True,
         cnn_in_channels=3,
         cnn_out_channels=16,
@@ -378,9 +381,7 @@ class ViT(nn.Module):
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         batch_norm_outputs=True,
-        initial_batch_norm=True,
-        linear_module=nn.Linear,
-        image_classes=100,
+        linear_module=SpectralNormLinear,
         head=SequencePoolClassificationHead,
     ):
 
@@ -404,6 +405,7 @@ class ViT(nn.Module):
 
         self.encoder = ViTEncoder(
             input_size=input_size,
+            initial_batch_norm=initial_batch_norm,
             cnn=cnn,
             cnn_in_channels=cnn_in_channels,
             cnn_out_channels=cnn_out_channels,
@@ -433,7 +435,6 @@ class ViT(nn.Module):
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             linear_module=linear_module,
-            initial_batch_norm=initial_batch_norm,
         )
 
         self.pool = head(

{broccoli_ml-0.13.0.dist-info → broccoli_ml-0.13.1.dist-info}/METADATA

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

{broccoli_ml-0.13.0.dist-info → broccoli_ml-0.13.1.dist-info}/RECORD

@@ -5,13 +5,13 @@ broccoli/assets/cifar100_eigenvectors_size_2.pt,sha256=DjXDOXMeuMpIqNuGhX9z-OWYV
 broccoli/assets/cifar100_eigenvectors_size_3.pt,sha256=gL6k0xtXYiYP6ZSvEiMBdJ7kIkT0AngTpDJHFQqwgxA,7173
 broccoli/cnn.py,sha256=jeRyKIAMWu1E3iyI14MGgSZuZivPMh12iqkqW9ilNjo,17785
 broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
-broccoli/linear.py,sha256=9ZwqC6kkgkr0uPoEjdi_Uq1QFHb4wCXzuU1r2pDreXM,2910
+broccoli/linear.py,sha256=jiGvLguxzkkmX14kRavaeg7IwN8jYJ06wn-NJ6Ivpzo,3008
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=E2JK5mQwJf75e23-JGcDoT7QxQf89DJReUo2et1LhRY,1716
-broccoli/transformer.py,sha256=E2GUsruxd5o3Eqrgra14IH0G0u_Wx4dIHhrUHrRqpLc,15929
+broccoli/transformer.py,sha256=GzkHlzCe4k2-ALMbKpQ0wdsOEKTap6gjOK-FiA7KP3k,15929
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=QNxdjB6xpW9G4hO7YknE5Uv7kXQYGbbdzEjqexQ5piY,16350
-broccoli_ml-0.13.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.13.0.dist-info/METADATA,sha256=Q99pX48E_VzsVBXvulm4TFhh3opWwfAVYl5-8pVr5HM,1257
-broccoli_ml-0.13.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.13.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=eEnb4hUwJUVymO3tD8V9JD-9i39ZkeNOYEDa9gwSL60,16398
+broccoli_ml-0.13.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.13.1.dist-info/METADATA,sha256=WWNGeC9F48atFNGqfw1Kv0i9QCGaiOIcHpIgAeMAAAw,1257
+broccoli_ml-0.13.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.13.1.dist-info/RECORD,,