Modify CCT Test Pretrained

Author: federicobrancasi

Tests/TestCCTPretrained.py

@@ -4,6 +4,8 @@
 #
 # Federico Brancasi <[email protected]>
 
+from typing import Tuple
+
 import brevitas.nn as qnn
 import pytest
 import torch
@@ -23,10 +25,61 @@
 from torch.utils.data import DataLoader, Subset
 from tqdm import tqdm
 
-from DeepQuant import brevitasToTrueQuant
+from DeepQuant.Transforms.Executor import TransformationExecutor
+from DeepQuant.Transforms.Transformations import LinearTransformation, MHATransformation
+from DeepQuant.Utils.ConsoleFormatter import ConsoleColor as cc
+from DeepQuant.Utils.CustomTracer import QuantTracer, customBrevitasTrace
+from DeepQuant.Utils.GraphPrinter import GraphModulePrinter
 from Tests.Models.CCT import cct_2_3x2_32
 
 
+def injectCustomForwards(
+    model: nn.Module,
+    exampleInput: torch.Tensor,
+    referenceOutput: torch.Tensor,
+    debug: bool = False,
+    checkEquivalence: bool = False,
+) -> Tuple[nn.Module, torch.Tensor]:
+    """Custom inject function for CCT that excludes ActivationTransformation."""
+    printer = GraphModulePrinter()
+
+    tracer = QuantTracer(debug=debug)
+
+    transformations = [
+        MHATransformation(),
+        LinearTransformation(),
+        # ActivationTransformation(),  # FBRANCASI: Commented out for CCT compatibility
+    ]
+
+    executor = TransformationExecutor(transformations, debug=debug, tracer=tracer)
+    transformedModel = executor.execute(model, exampleInput)
+
+    fxModel = customBrevitasTrace(
+        root=transformedModel,
+        tracer=tracer,
+    )
+    fxModel.recompile()
+
+    with torch.no_grad():
+        output = fxModel(exampleInput)
+
+    if checkEquivalence:
+        if torch.allclose(referenceOutput, output, atol=1e-5):
+            if debug:
+                print(cc.success("Injection of New Modules: output is consistent"))
+        else:
+            raise RuntimeError(
+                cc.error("Injection of New Modules changed the output significantly")
+            )
+
+    if debug:
+        print(cc.header("2. Network after Injection of New Modules"))
+        printer.printTabular(fxModel)
+        print()
+
+    return fxModel, output
+
+
 def evaluateModel(model, dataLoader, evalDevice, name="Model"):
     model.eval()
     correct = 0
@@ -133,8 +186,6 @@ def prepareFQCCT(model) -> nn.Module:
 
         quant_name = f"{node.name}_reshape_fix"
         model.add_module(quant_name, quant_identity)
-        # mark this QuantIdentity as “reshape fix”
-        quant_identity._is_reshape_fix = True
 
         with model.graph.inserting_after(node):
             quant_node = model.graph.call_module(quant_name, args=(node,))
@@ -265,7 +316,31 @@ def deepQuantTestCCT():
     FQAccuracy = evaluateModel(FQModel, valLoader, device, "FQ CCT-2")
 
     sampleInput = torch.randn(1, 3, 32, 32).to("cpu")
-    TQModel = brevitasToTrueQuant(FQModel, sampleInput, debug=True)
+    
+    # FBRANCASI: Override the injectCustomForwards function in the module before DeepQuant.Export imports it
+    import DeepQuant.Pipeline.Injection as injection_module
+
+    # FBRANCASI: Store original function
+    original_inject = injection_module.injectCustomForwards
+
+    # FBRANCASI: Override with our custom function
+    injection_module.injectCustomForwards = injectCustomForwards
+
+    # FBRANCASI: Force reload of Export module to pick up the override
+    import importlib
+
+    import DeepQuant.Export
+
+    importlib.reload(DeepQuant.Export)
+
+    try:
+        from DeepQuant.Export import brevitasToTrueQuant
+
+        TQModel = brevitasToTrueQuant(FQModel, sampleInput, debug=True)
+    finally:
+        # FBRANCASI: Restore original function and reload Export module again
+        injection_module.injectCustomForwards = original_inject
+        importlib.reload(DeepQuant.Export)
 
     numParameters = sum(p.numel() for p in TQModel.parameters())
     print(f"Number of parameters: {numParameters:,}")