Rename for better understanding

Author: federicobrancasi

DeepQuant/CustomForwards/Linear.py DeepQuant/CustomForwards/WBIOL.pyDeepQuant/CustomForwards/Linear.py renamed to DeepQuant/CustomForwards/WBIOL.py

@@ -9,7 +9,7 @@
 from torch import Tensor
 
 
-class WrapperLinear(nn.Module):
+class WrapperWBIOL(nn.Module):
     """Expose `inner_forward_impl` as a standalone submodule."""
 
     def __init__(self, innerForwardImpl: nn.Module) -> None:
@@ -22,7 +22,7 @@ def forward(
         return self.innerForwardImpl(quantInput, quantWeight, quantBias)
 
 
-def linearForward(self: QuantWeightBiasInputOutputLayer, inp: Tensor) -> Tensor:
+def WBIOLForward(self: QuantWeightBiasInputOutputLayer, inp: Tensor) -> Tensor:
     """Quant-in → quant-weight/bias → matmul → quant-out."""
     quantInput = self.input_quant(inp)
     quantWeight = self.weight_quant(self.weight)

DeepQuant/Pipeline/QuantSplit.py

@@ -9,7 +9,7 @@
 import torch
 import torch.nn as nn
 
-from DeepQuant.QuantManipulation.ParameterExtractor import (
+from DeepQuant.QuantManipulation.QuantizationParameterExtractor import (
     extractBrevitasProxyParams,
     printQuantParams,
 )

…/QuantManipulation/ParameterExtractor.py …lation/QuantizationParameterExtractor.pyDeepQuant/QuantManipulation/ParameterExtractor.py renamed to DeepQuant/QuantManipulation/QuantizationParameterExtractor.py DeepQuant/QuantManipulation/ParameterExtractor.py renamed to DeepQuant/QuantManipulation/QuantizationParameterExtractor.py

@@ -15,8 +15,8 @@
 from brevitas.proxy.runtime_quant import ActQuantProxyFromInjector
 
 
-def safeGetScale(quantObj: Any) -> Any:
-    """Safely extract scale parameter from quantization object."""
+def getScale(quantObj: Any) -> Any:
+    """Extract scale parameter from quantization object."""
     if quantObj is None:
         return None
     maybeScale = quantObj.scale() if callable(quantObj.scale) else quantObj.scale
@@ -32,8 +32,8 @@ def safeGetScale(quantObj: Any) -> Any:
         return None
 
 
-def safeGetZeroPoint(quantObj: Any) -> Any:
-    """Safely extract zero point parameter from quantization object."""
+def getZeroPoint(quantObj: Any) -> Any:
+    """Extract zero point parameter from quantization object."""
     if quantObj is None:
         return None
     maybeZp = (
@@ -51,16 +51,16 @@ def safeGetZeroPoint(quantObj: Any) -> Any:
         return None
 
 
-def safeGetIsSigned(quantObj: Any) -> bool:
-    """Safely determine if quantization is signed."""
+def getIsSigned(quantObj: Any) -> bool:
+    """Determine if quantization is signed."""
     if hasattr(quantObj, "is_signed"):
         return getattr(quantObj, "is_signed")
     if hasattr(quantObj, "min_val"):
         try:
             return quantObj.min_val < 0
         except Exception:
             pass
-    zp = safeGetZeroPoint(quantObj)
+    zp = getZeroPoint(quantObj)
     if zp is not None:
         # If zero_point is near zero, assume unsigned quantization.
         return not (abs(zp) < 1e-5)
@@ -82,10 +82,10 @@ def recurseModules(parentMod: nn.Module, prefix: str = "") -> None:
                     BiasQuantProxyFromInjector,
                 ),
             ):
-                scl = safeGetScale(childMod)
-                zp = safeGetZeroPoint(childMod)
+                scl = getScale(childMod)
+                zp = getZeroPoint(childMod)
                 bw = childMod.bit_width()
-                isSigned = safeGetIsSigned(childMod)
+                isSigned = getIsSigned(childMod)
                 paramsDict[fullName] = {
                     "scale": scl,
                     "zero_point": zp,

DeepQuant/Transforms/Transformations.py

@@ -14,7 +14,7 @@
 from brevitas.nn.quant_mha import QuantMultiheadAttention
 
 from DeepQuant.CustomForwards.Activations import WrapperActivation, activationForward
-from DeepQuant.CustomForwards.Linear import WrapperLinear, linearForward
+from DeepQuant.CustomForwards.WBIOL import WBIOLForward, WrapperWBIOL
 from DeepQuant.CustomForwards.MultiHeadAttention import mhaForward
 from DeepQuant.Transforms.Base import TransformationPass
 from DeepQuant.Utils.CustomTracer import QuantTracer
@@ -33,11 +33,11 @@ def injectForward(
         self, module: nn.Module, tracer: Optional[QuantTracer] = None
     ) -> None:
         """Inject custom forward for linear layers."""
-        module.wrappedInnerForwardImpl = WrapperLinear(module.inner_forward_impl)
-        module.forward = linearForward.__get__(module)
+        module.wrappedInnerForwardImpl = WrapperWBIOL(module.inner_forward_impl)
+        module.forward = WBIOLForward.__get__(module)
 
         if tracer:
-            tracer.registerLeafModule(WrapperLinear)
+            tracer.registerLeafModule(WrapperWBIOL)
             tracer.registerNonLeafModule(QuantWeightBiasInputOutputLayer)
 
 

Tests/TestYOLOv5.py

@@ -23,7 +23,7 @@ def prepareYOLOv5Backbone() -> nn.Module:
     """Prepare a quantized partial YOLOv5 model for testing."""
     from ultralytics import YOLO
 
-    model = YOLO("Models/yolov5n.pt")
+    model = YOLO("Models/yolov5nu.pt")
     pytorchModel = model.model
 
     # FBRANCASI: Just first few layers for simplicity