First round of refactoring

Author: Victor-Jung

.github/workflows/CI.yml

@@ -36,9 +36,9 @@ jobs:
           pip install -e .
       - name: Run Tests
         run: |
-          pytest Tests/test_simple_cnn.py
-          pytest Tests/test_simple_mha.py
-          pytest Tests/test_simple_nn.py
+          pytest Tests/TestSimpleCNN.py
+          pytest Tests/TestSimpleMHA.py
+          pytest Tests/TestSimpleNN.py
 
   model-tests:
     runs-on: ubuntu-latest
@@ -55,4 +55,4 @@ jobs:
           pip install -e .
       - name: Run Tests
         run: |
-          pytest Tests/test_mnist.py
+          pytest Tests/TestMnist.py

DeepQuant/custom_forwards/activations.py DeepQuant/CustomForwards/Activations.pyDeepQuant/custom_forwards/activations.py renamed to DeepQuant/CustomForwards/Activations.py

@@ -4,11 +4,7 @@
 #
 # Federico Brancasi <[email protected]>
 
-"""
-Custom forward implementations for Brevitas QuantActivation layers.
-"""
 
-import torch
 import torch.nn as nn
 from torch import Tensor
 from brevitas.nn.quant_layer import QuantNonLinearActLayer
@@ -22,15 +18,15 @@ class InnerForwardImplWrapperActivation(nn.Module):
     so that FX tracing can display it as a separate node.
     """
 
-    def __init__(self, act_impl: nn.Module) -> None:
+    def __init__(self, actImpl: nn.Module) -> None:
         """
         Args:
             act_impl: The original activation function module (e.g. an instance of nn.ReLU).
         """
         super().__init__()
-        self.act_impl = act_impl
+        self.actImpl = actImpl
 
-    def forward(self, quant_input: Tensor) -> Tensor:
+    def forward(self, quantInput: Tensor) -> Tensor:
         """
         Applies the wrapped activation function.
 
@@ -40,10 +36,10 @@ def forward(self, quant_input: Tensor) -> Tensor:
         Returns:
             Output tensor after applying the activation.
         """
-        return self.act_impl(quant_input)
+        return self.actImpl(quantInput)
 
 
-def quant_activation_forward(self: QuantNonLinearActLayer, inp: Tensor) -> Tensor:
+def quantActivationForward(self: QuantNonLinearActLayer, inp: Tensor) -> Tensor:
     """
     Unrolled forward pass for a Brevitas QuantActivation layer.
 
@@ -59,11 +55,12 @@ def quant_activation_forward(self: QuantNonLinearActLayer, inp: Tensor) -> Tenso
     Returns:
         Output tensor after applying activation and output quantization.
     """
-    quant_input = self.input_quant(inp) if self.input_quant is not None else inp
+    quantInput = self.input_quant(inp) if self.input_quant is not None else inp
     # Use the wrapped activation if available; otherwise pass through.
-    if hasattr(self, "wrapped_act_impl"):
-        output = self.wrapped_act_impl(quant_input)
+    if hasattr(self, "wrappedActImpl"):
+        output = self.wrappedActImpl(quantInput)
     else:
-        output = quant_input
-    quant_output = self.act_quant(output) if self.act_quant is not None else output
-    return quant_output
+        output = quantInput
+        import IPython; IPython.embed()
+    quantOutput = self.act_quant(output) if self.act_quant is not None else output
+    return quantOutput

DeepQuant/custom_forwards/linear.py DeepQuant/CustomForwards/Linear.pyDeepQuant/custom_forwards/linear.py renamed to DeepQuant/CustomForwards/Linear.py

@@ -4,39 +4,35 @@
 #
 # Federico Brancasi <[email protected]>
 
-"""
-Custom forward implementations for Brevitas QuantLinear layers.
-"""
 
-import torch
 import torch.nn as nn
 from torch import Tensor
 from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer
 
 
 class InnerForwardImplWrapperLinear(nn.Module):
     """
-    A small wrapper around the 'inner_forward_impl' of a Brevitas QuantLinear
+    A small wrapper around the 'innerForwardImpl' of a Brevitas QuantLinear
     (QuantWeightBiasInputOutputLayer).
 
-    We want to expose the logic within 'inner_forward_impl' as a standalone
+    We want to expose the logic within 'innerForwardImpl' as a standalone
     submodule, so that FX tracing can see it as a leaf.
     """
 
-    def __init__(self, inner_forward_impl: nn.Module) -> None:
+    def __init__(self, innerForwardImpl: nn.Module) -> None:
         """
         Args:
-            inner_forward_impl: The original function that processes
+            innerForwardImpl: The original function that processes
                                 (quant_input, quant_weight, quant_bias).
         """
         super().__init__()
-        self.inner_forward_impl = inner_forward_impl
+        self.innerForwardImpl = innerForwardImpl
 
     def forward(
-        self, quant_input: Tensor, quant_weight: Tensor, quant_bias: Tensor
+        self, quantInput: Tensor, quantWeight: Tensor, quantBias: Tensor
     ) -> Tensor:
         """
-        Applies the wrapped inner_forward_impl.
+        Applies the wrapped innerForwardImpl.
 
         Args:
             quant_input: Input after input_quant.
@@ -46,18 +42,18 @@ def forward(
         Returns:
             A torch.Tensor with the linear operation applied.
         """
-        return self.inner_forward_impl(quant_input, quant_weight, quant_bias)
+        return self.innerForwardImpl(quantInput, quantWeight, quantBias)
 
 
-def quantWBIOL_forward(self: QuantWeightBiasInputOutputLayer, inp: Tensor) -> Tensor:
+def quantWBIOLForward(self: QuantWeightBiasInputOutputLayer, inp: Tensor) -> Tensor:
     """
     Unrolled forward pass for a Brevitas QuantLinear:
 
     Steps:
       1) self.input_quant
       2) self.weight_quant
       3) self.bias_quant (if bias is present)
-      4) inner_forward_impl (wrapped)
+      4) innerForwardImpl (wrapped)
       5) self.output_quant
 
     Args:
@@ -67,13 +63,13 @@ def quantWBIOL_forward(self: QuantWeightBiasInputOutputLayer, inp: Tensor) -> Te
     Returns:
         Output Tensor after the unrolled quantized linear steps.
     """
-    quant_input = self.input_quant(inp)
-    quant_weight = self.weight_quant(self.weight)
+    quantInput = self.input_quant(inp)
+    quantWeight = self.weight_quant(self.weight)
 
-    quant_bias = None
+    quantBias = None
     if self.bias is not None:
-        quant_bias = self.bias_quant(self.bias, quant_input, quant_weight)
+        quantBias = self.bias_quant(self.bias, quantInput, quantWeight)
 
-    output = self.wrapped_inner_forward_impl(quant_input, quant_weight, quant_bias)
-    quant_output = self.output_quant(output)
-    return quant_output
+    output = self.wrappedInnerForwardImpl(quantInput, quantWeight, quantBias)
+    quantOutput = self.output_quant(output)
+    return quantOutput

…nt/custom_forwards/multiheadattention.py …ant/CustomForwards/MultiHeadAttention.pyDeepQuant/custom_forwards/multiheadattention.py renamed to DeepQuant/CustomForwards/MultiHeadAttention.py DeepQuant/custom_forwards/multiheadattention.py renamed to DeepQuant/CustomForwards/MultiHeadAttention.py

@@ -4,9 +4,6 @@
 #
 # Federico Brancasi <[email protected]>
 
-"""
-Custom forward implementation for Brevitas QuantMultiheadAttention.
-"""
 
 import math
 import torch
@@ -15,7 +12,7 @@
 from brevitas.nn.quant_mha import QuantMultiheadAttention
 
 
-def unrolled_quant_mha_forward(
+def unrolledQuantMhaForward(
     self: QuantMultiheadAttention, query: Tensor, key: Tensor, value: Tensor
 ) -> Tensor:
     """
@@ -39,52 +36,52 @@ def unrolled_quant_mha_forward(
         after the unrolled MHA steps.
     """
     # 1) Q, K, V projections
-    q_out = self.q_proj(query)
-    k_out = self.k_proj(key)
-    v_out = self.v_proj(value)
+    qOut = self.q_proj(query)
+    kOut = self.k_proj(key)
+    vOut = self.v_proj(value)
 
     # 2) Multi-head reshape
-    seq_len, batch_size, embed_dim = q_out.shape
-    head_dim = embed_dim // self.num_heads
+    seqLen, batchSize, embedDim = qOut.shape
+    headDim = embedDim // self.num_heads
 
-    q_out = (
-        q_out.view(seq_len, batch_size, self.num_heads, head_dim)
+    qOut = (
+        qOut.view(seqLen, batchSize, self.num_heads, headDim)
         .permute(1, 2, 0, 3)
-        .reshape(batch_size * self.num_heads, seq_len, head_dim)
+        .reshape(batchSize * self.num_heads, seqLen, headDim)
     )
-    k_out = (
-        k_out.view(seq_len, batch_size, self.num_heads, head_dim)
+    kOut = (
+        kOut.view(seqLen, batchSize, self.num_heads, headDim)
         .permute(1, 2, 0, 3)
-        .reshape(batch_size * self.num_heads, seq_len, head_dim)
+        .reshape(batchSize * self.num_heads, seqLen, headDim)
     )
-    v_out = (
-        v_out.view(seq_len, batch_size, self.num_heads, head_dim)
+    vOut = (
+        vOut.view(seqLen, batchSize, self.num_heads, headDim)
         .permute(1, 2, 0, 3)
-        .reshape(batch_size * self.num_heads, seq_len, head_dim)
+        .reshape(batchSize * self.num_heads, seqLen, headDim)
     )
 
     # 3) Scale queries, then quantize
-    q_scaled = q_out / math.sqrt(head_dim)
-    q_scaled = self.q_scaled_quant(q_scaled)
+    qScaled = qOut / math.sqrt(headDim)
+    qScaled = self.q_scaled_quant(qScaled)
 
     # 4) Transpose + quantize K, compute attention weights
-    k_t = k_out.transpose(-2, -1)
+    k_t = kOut.transpose(-2, -1)
     k_t = self.k_transposed_quant(k_t)
 
-    attn_weights = torch.bmm(q_scaled, k_t)
-    attn_weights = self.softmax_input_quant(attn_weights)
-    attn_weights = F.softmax(attn_weights, dim=-1)
-    attn_weights = self.attn_output_weights_quant(attn_weights)
+    attnWeights = torch.bmm(qScaled, k_t)
+    attnWeights = self.softmax_input_quant(attnWeights)
+    attnWeights = F.softmax(attnWeights, dim=-1)
+    attnWeights = self.attn_output_weights_quant(attnWeights)
 
     # 5) Quantize V, multiply, reshape back, and final out projection
-    v_out = self.v_quant(v_out)
-    attn_output = torch.bmm(attn_weights, v_out)
+    vOut = self.v_quant(vOut)
+    attnOutput = torch.bmm(attnWeights, vOut)
 
-    attn_output = (
-        attn_output.view(batch_size, self.num_heads, seq_len, head_dim)
+    attnOutput = (
+        attnOutput.view(batchSize, self.num_heads, seqLen, headDim)
         .permute(2, 0, 1, 3)
-        .reshape(seq_len, batch_size, embed_dim)
+        .reshape(seqLen, batchSize, embedDim)
     )
 
-    attn_output = self.out_proj(attn_output)
-    return attn_output
+    attnOutput = self.out_proj(attnOutput)
+    return attnOutput

DeepQuant/custom_tracer.py DeepQuant/CustomTracer.pyDeepQuant/custom_tracer.py renamed to DeepQuant/CustomTracer.py

@@ -9,7 +9,6 @@
 """
 
 import torch.nn as nn
-import torch.fx as fx
 from brevitas.fx.brevitas_tracer import (
     _symbolic_trace,
     _is_brevitas_leaf_module,
@@ -30,8 +29,8 @@ class CustomBrevitasTracer(Tracer):
 
     def __init__(
         self,
-        leaf_classes: Optional[List[Type[nn.Module]]] = None,
-        non_leaf_classes: Optional[List[Type[nn.Module]]] = None,
+        leafClasses: Optional[List[Type[nn.Module]]] = None,
+        nonLeafClasses: Optional[List[Type[nn.Module]]] = None,
         debug: bool = False,
     ) -> None:
         """
@@ -43,31 +42,31 @@ def __init__(
             debug: Whether to print debug information during tracing.
         """
         super().__init__()
-        self.leaf_classes = leaf_classes if leaf_classes is not None else []
-        self.non_leaf_classes = non_leaf_classes if non_leaf_classes is not None else []
+        self.leafClasses = leafClasses if leafClasses is not None else []
+        self.nonLeafClasses = nonLeafClasses if nonLeafClasses is not None else []
         self.debug = debug
 
-    def register_leaf_module(self, module_cls: Type[nn.Module]) -> None:
+    def registerLeafModule(self, moduleCls: Type[nn.Module]) -> None:
         """
         Add a module class to the list of leaf modules.
 
         Args:
             module_cls: The module class to register as a leaf module.
         """
-        if module_cls not in self.leaf_classes:
-            self.leaf_classes.append(module_cls)
+        if moduleCls not in self.leafClasses:
+            self.leafClasses.append(moduleCls)
 
-    def register_non_leaf_module(self, module_cls: Type[nn.Module]) -> None:
+    def registerNonLeafModule(self, moduleCls: Type[nn.Module]) -> None:
         """
         Add a module class to the list of non-leaf modules.
 
         Args:
             module_cls: The module class to register as a non-leaf module.
         """
-        if module_cls not in self.non_leaf_classes:
-            self.non_leaf_classes.append(module_cls)
+        if moduleCls not in self.nonLeafClasses:
+            self.nonLeafClasses.append(moduleCls)
 
-    def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:
+    def is_leaf_module(self, m: nn.Module, moduleQualifiedName: str) -> bool:
         """
         Determine whether a module should be treated as a leaf module.
 
@@ -84,16 +83,16 @@ def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:
             bool: True if the module should be treated as a leaf module, False otherwise.
         """
         # First check explicitly registered classes
-        if any(isinstance(m, lc) for lc in self.leaf_classes):
+        if any(isinstance(m, lc) for lc in self.leafClasses):
             return True
-        if any(isinstance(m, nlc) for nlc in self.non_leaf_classes):
+        if any(isinstance(m, nlc) for nlc in self.nonLeafClasses):
             return False
         # Fall back to default Brevitas behavior
-        return _is_brevitas_leaf_module(m, module_qualified_name)
+        return _is_brevitas_leaf_module(m, moduleQualifiedName)
 
 
-def custom_brevitas_trace(
-    root: nn.Module, concrete_args=None, tracer: Optional[CustomBrevitasTracer] = None
+def customBrevitasTrace(
+    root: nn.Module, concreteArgs=None, tracer: Optional[CustomBrevitasTracer] = None
 ) -> GraphModule:
     """
     Create an FX GraphModule using the CustomBrevitasTracer.
@@ -108,4 +107,4 @@ def custom_brevitas_trace(
     """
     if tracer is None:
         tracer = CustomBrevitasTracer()
-    return _symbolic_trace(tracer, root, concrete_args)
+    return _symbolic_trace(tracer, root, concreteArgs)