update the unit test

aidge_quantization/unit_tests/test_ptq.py

 import unittest
-import gzip
 import numpy as np
-from pathlib import Path
-
+import gzip
 import aidge_core
-import aidge_backend_cpu
 import aidge_onnx
+import aidge_backend_cpu
 import aidge_quantization
+import sys
+from pathlib import Path
 
-from aidge_core import Log, Level
-
+"""
+    Unit test for the PTQ pipeline
+    ==============================
+    This script is designed to test and validate the accuracy of five small topologies on the MNIST dataset :
+    ["MiniResNet", "ConvNet", "BranchNetV4", "TestNet", "MLP"]
+    It compares the results for three configurations : baseline, quantization, and quantization with single shift. 
+    The value of delta represents the tolerance of the tests.
+"""
+aidge_core.Log.set_console_level(aidge_core.Level.Error)  # Reduce useless logs
 # --------------------------------------------------------------
-# CONFIGS
+# CONFIGURATION
 # --------------------------------------------------------------
 
-NB_SAMPLES      = 1000               # max : 1000
-SAMPLE_SHAPE    = (1, 1, 28, 28)
-MODEL_NAME      = 'MiniResNet.onnx'  # 'ConvNet.onnx'
-ACCURACIES      = (95.4, 94.4)       # (97.9, 97.7)
-NB_BITS         = 4
+NB_SAMPLES   = 1000
+SAMPLE_SHAPE = (1, 1, 28, 28)
+NB_BITS      = 4
+TARGET_TYPE  = aidge_core.dtype.int32
+CLIPPING     = aidge_quantization.Clipping.MSE
+NO_QUANT     = False
+OPTIM_SIGNS  = True
+FOLD_GRAPH   = True
+DELTA        = 0.05
+
+EXPECTED_RESULTS = {
+    "MiniResNet.onnx"  : (95.4, 94.4, 95.0),
+    "ConvNet.onnx"     : (97.9, 97.2, 96.7),
+    "BranchNetV4.onnx" : (93.8, 92.7, 93.7),
+    "TestNet.onnx"     : (95.5, 94.0, 94.5),
+    "MLP.onnx"         : (94.7, 92.9, 93.8)
+}
 
 # --------------------------------------------------------------
 # UTILS
 # --------------------------------------------------------------
 
 def propagate(model, scheduler, sample):
+    sample = np.reshape(sample, SAMPLE_SHAPE)
     input_tensor = aidge_core.Tensor(sample)
     scheduler.forward(True, [input_tensor])
     output_node = model.get_output_nodes().pop()
     output_tensor = output_node.get_operator().get_output(0)
     return np.array(output_tensor)
 
-def prepare_sample(sample):
-    sample = np.reshape(sample, SAMPLE_SHAPE)
-    return sample.astype('float32')
-
 def compute_accuracy(model, samples, labels):
-    acc = 0
     scheduler = aidge_core.SequentialScheduler(model)
-    for i, sample in enumerate(samples):
-        x = prepare_sample(sample)
-        y = propagate(model, scheduler, x)
-        if labels[i] == np.argmax(y):
-            acc += 1
+    acc = sum(labels[i] == np.argmax(propagate(model, scheduler, x)) for i, x in enumerate(samples))
     return acc / len(samples)
 
 # --------------------------------------------------------------
 # TEST CLASS
 # --------------------------------------------------------------
 
-class test_ptq(unittest.TestCase):
+class TestQuantization(unittest.TestCase):
 
     def setUp(self):
-
-        # load the samples / labels (numpy)
-
         curr_file_dir = Path(__file__).parent.resolve()
         self.samples = np.load(gzip.GzipFile(curr_file_dir / 'assets/mnist_samples.npy.gz', "r"))
-        self.labels  = np.load(gzip.GzipFile(curr_file_dir / 'assets/mnist_labels.npy.gz',  "r"))
-
-        # load the model in AIDGE
+        self.labels = np.load(gzip.GzipFile(curr_file_dir / 'assets/mnist_labels.npy.gz', "r"))
+        self.quant_samples = np.round(self.samples.copy() * (2**(NB_BITS-1)-1))
 
-        self.model = aidge_onnx.load_onnx(curr_file_dir / "assets/" / MODEL_NAME, verbose=False)
-        aidge_core.remove_flatten(self.model)
+    def run_model_test(self, model_name):
 
-        self.model.set_datatype(aidge_core.dtype.float32)
-        self.model.set_backend("cpu")
+        expected_base, expected_quant, expected_quant_ssa = EXPECTED_RESULTS[model_name]
 
-    def tearDown(self):
-        pass
+        # load the model ...
 
+        model_path = Path(__file__).parent / "assets" / model_name
+        model = aidge_onnx.load_onnx(model_path, verbose=False)
+        aidge_core.remove_flatten(model)
 
-    def test_model(self):
+        model.set_datatype(aidge_core.dtype.float32)
+        model.set_backend("cpu")
+        
+        # create the tensor subset
 
-        Log.set_console_level(Level.Info)
-        # compute the base accuracy
-        accuracy = compute_accuracy(self.model, self.samples[0:NB_SAMPLES], self.labels)
-        self.assertAlmostEqual(accuracy * 100, ACCURACIES[0], msg='base accuracy does not meet the baseline !', delta=0.1)
+        tensors = [aidge_core.Tensor(np.reshape(sample, SAMPLE_SHAPE)) for sample in self.samples[:NB_SAMPLES]]
 
-    def test_quant_model(self):
+        # BASELINE ACCURACY
 
-        Log.set_console_level(Level.Debug)
+        base_accuracy = compute_accuracy(model, tensors, self.labels[:NB_SAMPLES])
+        self.assertAlmostEqual(base_accuracy * 100, expected_base, delta=DELTA, msg=f"[X] Baseline accuracy mismatch for {model_name}. Expected accuracy was: {expected_base}, but got: {base_accuracy * 100}")
 
-        # create the calibration dataset
-
-        tensors = []
-        for sample in self.samples[0:NB_SAMPLES]:
-            sample = prepare_sample(sample)
-            tensor = aidge_core.Tensor(sample)
-            tensors.append(tensor)
-
-        # quantize the model
+        # QUANTIZED ACCURACY
 
         aidge_quantization.quantize_network(
-            self.model,
-            NB_BITS,
-            tensors,
-            clipping_mode=aidge_quantization.Clipping.MSE,
-            no_quantization=False,
-            optimize_signs=True,
-            single_shift=False
-        )
-
-        # rescale the inputs
-
-        scaling = 2**(NB_BITS-1)-1
-        for i in range(NB_SAMPLES):
-            self.samples[i] = self.samples[i]*scaling # XXX np.round ???
-
-        # compute the quantized accuracy
-
-        accuracy = compute_accuracy(self.model, self.samples, self.labels)
-        self.assertAlmostEqual(accuracy * 100, ACCURACIES[1], msg='quantized accuracy does not meet the baseline !', delta=0.1)
-
+            network=model,
+            nb_bits=NB_BITS,
+            calibration_set=tensors,
+            target_type=TARGET_TYPE,
+            clipping_mode=CLIPPING,
+            no_quant=NO_QUANT,
+            optimize_signs=OPTIM_SIGNS,
+            single_shift=False,
+            use_cuda=False,
+            fold_graph=FOLD_GRAPH)
+
+        quant_accuracy = compute_accuracy(model, self.quant_samples[:NB_SAMPLES], self.labels)
+        self.assertAlmostEqual(quant_accuracy * 100, expected_quant, delta=DELTA, msg=f"[X] Quantized accuracy mismatch for {model_name}. Expected accuracy was: {expected_quant}, but got: {quant_accuracy * 100}")
+        
+        # QUANTIZED ACCURACY WITH SSA
+
+        model = aidge_onnx.load_onnx(model_path, verbose=False)
+        model.set_datatype(aidge_core.dtype.float32)
+        model.set_backend("cpu")
+        
+        aidge_quantization.quantize_network(
+            network=model,
+            nb_bits=NB_BITS,
+            calibration_set=tensors,
+            target_type=TARGET_TYPE,
+            clipping_mode=CLIPPING,
+            no_quant=NO_QUANT,
+            optimize_signs=OPTIM_SIGNS,
+            single_shift=True,
+            use_cuda=False,
+            fold_graph=FOLD_GRAPH)
+
+        quant_accuracy_ssa = compute_accuracy(model, self.quant_samples[:NB_SAMPLES], self.labels)
+        self.assertAlmostEqual(quant_accuracy_ssa * 100, expected_quant_ssa, delta=DELTA, msg=f"[X] Quantized accuracy (with SSA) mismatch for {model_name}. Expected accuracy was: {expected_quant_ssa}, but got: {quant_accuracy_ssa * 100}")
+
+    def test_models(self):
+        for model in EXPECTED_RESULTS.keys():
+            with self.subTest(model=model):
+                self.run_model_test(model)
 
 if __name__ == '__main__':
     unittest.main()