Codebase Refactor

Author: federicobrancasi

DeepQuant/Export.py

@@ -299,7 +299,7 @@ def exportQuantModel(
         # f=EXPORT_FOLDER / "4_model_dequant_moved.onnx",
         f=onnxFile,
         opset_version=13,
-        keep_initializers_as_inputs=True,
+        keep_initializers_as_inputs=False,  # FBRANCASI: This prevent the onnx warnings
         do_constant_folding=False,
         input_names=["input"],
         output_names=["output"],

Tests/TestConv.py

@@ -5,7 +5,6 @@
 # Victor Jung <[email protected]>
 # Federico Brancasi <[email protected]>
 
-
 import pytest
 import torch
 import torch.nn as nn
@@ -19,8 +18,9 @@
 
 
 class QuantConvNet(nn.Module):
+    """Simple quantized CNN with a single conv layer."""
 
-    convAndLinQuantParams = {
+    convQuantParams = {
         "bias": True,
         "weight_bit_width": 4,
         "bias_quant": Int32Bias,
@@ -30,31 +30,26 @@ class QuantConvNet(nn.Module):
         "return_quant_tensor": True,
     }
 
-    def __init__(self, in_channels: int = 1) -> None:
+    def __init__(self, inChannels: int = 1) -> None:
         super().__init__()
         self.inputQuant = qnn.QuantIdentity(return_quant_tensor=True)
-
         self.conv1 = qnn.QuantConv2d(
-            in_channels=in_channels,
+            in_channels=inChannels,
             out_channels=16,
             kernel_size=3,
             padding=1,
-            **QuantConvNet.convAndLinQuantParams,
+            **QuantConvNet.convQuantParams,
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-
         x = self.inputQuant(x)
         x = self.conv1(x)
-
         return x
 
 
 @pytest.mark.SingleLayerTests
 def deepQuantTestConv() -> None:
-
     torch.manual_seed(42)
-
     model = QuantConvNet().eval()
     sampleInput = torch.randn(1, 1, 28, 28)
     exportQuantModel(model, sampleInput, debug=True)

Tests/TestLinear.py

@@ -4,14 +4,9 @@
 #
 # Federico Brancasi <[email protected]>
 
-
 import pytest
-
-### PyTorch Imports ###
 import torch
 import torch.nn as nn
-
-### Brevitas Import ###
 import brevitas.nn as qnn
 from brevitas.quant.scaled_int import (
     Int8ActPerTensorFloat,
@@ -22,15 +17,14 @@
 
 
 class QuantLinearNet(nn.Module):
+    """Simple quantized network with a single linear layer."""
 
-    def __init__(self, in_features: int = 16, hidden_features: int = 32) -> None:
+    def __init__(self, inFeatures: int = 16, hiddenFeatures: int = 32) -> None:
         super().__init__()
-
         self.inputQuant = qnn.QuantIdentity(return_quant_tensor=True)
-
         self.linear1 = qnn.QuantLinear(
-            in_features=in_features,
-            out_features=hidden_features,
+            in_features=inFeatures,
+            out_features=hiddenFeatures,
             bias=True,
             weight_bit_width=4,
             bias_quant=Int32Bias,
@@ -41,19 +35,14 @@ def __init__(self, in_features: int = 16, hidden_features: int = 32) -> None:
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-
         x = self.inputQuant(x)
         x = self.linear1(x)
-
         return x
 
 
 @pytest.mark.SingleLayerTests
 def deepQuantTestLinear() -> None:
-    
     torch.manual_seed(42)
-
     model = QuantLinearNet().eval()
     sampleInput = torch.randn(1, 4, 16)
-
     exportQuantModel(model, sampleInput, debug=True)

Tests/TestMHSA.py

@@ -4,7 +4,6 @@
 #
 # Federico Brancasi <[email protected]>
 
-
 import pytest
 import torch
 import torch.nn as nn
@@ -21,22 +20,18 @@
 
 
 class QuantMHSANet(nn.Module):
+    """Simple quantized network with multi-head self-attention."""
 
-    def __init__(self, embed_dim: int, num_heads: int) -> None:
-        """
-        Args:
-            embed_dim: The dimension of each embedding vector.
-            num_heads: The number of attention heads.
-        """
+    def __init__(self, embedDim: int, numHeads: int) -> None:
         super().__init__()
         self.inputQuant = qnn.QuantIdentity(return_quant_tensor=True)
         self.mha = qnn.QuantMultiheadAttention(
-            embed_dim=embed_dim,
-            num_heads=num_heads,
+            embed_dim=embedDim,
+            num_heads=numHeads,
             dropout=0.0,
             bias=True,
-            packed_in_proj=False,  # separate Q, K, V
-            batch_first=False,  # expects (sequence, batch, embed_dim)
+            packed_in_proj=False,  # FBRANCASI: separate Q, K, V
+            batch_first=False,  # FBRANCASI: expects (sequence, batch, embed_dim)
             in_proj_input_quant=Int8ActPerTensorFloat,
             in_proj_weight_quant=Int8WeightPerTensorFloat,
             in_proj_bias_quant=Int32Bias,
@@ -51,27 +46,14 @@ def __init__(self, embed_dim: int, num_heads: int) -> None:
         )
 
     def forward(self, x: Tensor) -> Tensor:
-        """
-        Forward pass that first quantizes the input, then applies multi-head attention.
-
-        Args:
-            x: Input tensor of shape [sequence_len, batch_size, embed_dim].
-
-        Returns:
-            A tuple (output, None) as per the Brevitas MHA API, where output has shape
-            [sequence_len, batch_size, embed_dim].
-        """
         x = self.inputQuant(x)
         out = self.mha(x, x, x)
         return out
 
 
 @pytest.mark.SingleLayerTests
 def deepQuantTestMHSA() -> None:
-
     torch.manual_seed(42)
-
-    model = QuantMHSANet(embed_dim=16, num_heads=4).eval()
+    model = QuantMHSANet(embedDim=16, numHeads=4).eval()
     sampleInput = torch.randn(10, 2, 16)
-
     exportQuantModel(model, sampleInput)

Tests/TestMobileNetV3Small.py

@@ -2,7 +2,7 @@
 # Licensed under the Apache License, Version 2.0, see LICENSE for details.
 # SPDX-License-Identifier: Apache-2.0
 #
-# Victor Juing <[email protected]>
+# Victor Jung <[email protected]>
 
 import pytest
 import torch
@@ -23,18 +23,10 @@
 
 
 def prepareMBNetV3Model() -> nn.Module:
-    """
-    Prepare a quantized MobileNetV3Small model for testing.
-    Steps:
-      1) Load the torchvision MobileNetV3Small.
-      2) Convert it to eval mode.
-      3) Preprocess and adapt average pooling.
-      4) Quantize it using Brevitas.
-
-    Returns:
-        A quantized MobileNetV3Small model ready for export tests.
-    """
-    baseModel = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)
+    """Prepare a quantized MobileNetV3Small model for testing."""
+    baseModel = models.mobilenet_v3_small(
+        weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1
+    )
     baseModel = baseModel.eval()
 
     computeLayerMap = {
@@ -99,9 +91,7 @@ def prepareMBNetV3Model() -> nn.Module:
     baseModel = preprocess_for_quantize(
         baseModel, equalize_iters=20, equalize_scale_computation="range"
     )
-    baseModel = AdaptiveAvgPoolToAvgPool().apply(
-        baseModel, torch.ones(1, 3, 224, 224)
-    )
+    baseModel = AdaptiveAvgPoolToAvgPool().apply(baseModel, torch.ones(1, 3, 224, 224))
 
     quantizedModel = quantize(
         graph_model=baseModel,
@@ -115,10 +105,7 @@ def prepareMBNetV3Model() -> nn.Module:
 
 @pytest.mark.ModelTests
 def deepQuantTestMobileNetV3Small() -> None:
-
     torch.manual_seed(42)
-
-    quantizedModel = prepareMBNetV3Model()
+    model = prepareMBNetV3Model()
     sampleInput = torch.randn(1, 3, 224, 224)
-
-    exportQuantModel(quantizedModel, sampleInput, debug=True)
+    exportQuantModel(model, sampleInput, debug=True)