feat: update compile_torch_model to return compiled quantized module

closes #898

tests/quantization/test_compilation.py

@@ -85,7 +85,7 @@ def test_quantized_module_compilation(
         homomorphic_predictions = homomorphic_predictions.reshape(dequant_predictions.shape)
 
         # Make sure homomorphic_predictions are the same as dequant_predictions
-        if numpy.isclose(homomorphic_predictions.ravel(), dequant_predictions.ravel()).all():
+        if numpy.isclose(homomorphic_predictions, dequant_predictions).all():
             return
 
         # Bad computation after nb_tries

tests/torch/test_compile_torch.py

@@ -1,62 +1,79 @@
 """Tests for the torch to numpy module."""
+import numpy
 import pytest
-import torch
 from torch import nn
 
+from concrete.quantization import QuantizedArray
 from concrete.torch.compile import compile_torch_model
 
+# INPUT_OUTPUT_FEATURE is the number of input and output of each of the network layers.
+# (as well as the input of the network itself)
+INPUT_OUTPUT_FEATURE = [1, 2, 3]
+
 
 class FC(nn.Module):
     """Torch model for the tests"""
 
-    def __init__(self):
+    def __init__(self, input_output):
         super().__init__()
-        self.fc1 = nn.Linear(in_features=32 * 32 * 3, out_features=128)
+        self.fc1 = nn.Linear(in_features=input_output, out_features=input_output)
         self.sigmoid1 = nn.Sigmoid()
-        self.fc2 = nn.Linear(in_features=128, out_features=64)
-        self.sigmoid2 = nn.Sigmoid()
-        self.fc3 = nn.Linear(in_features=64, out_features=64)
-        self.sigmoid3 = nn.Sigmoid()
-        self.fc4 = nn.Linear(in_features=64, out_features=64)
-        self.sigmoid4 = nn.Sigmoid()
-        self.fc5 = nn.Linear(in_features=64, out_features=10)
+        self.fc2 = nn.Linear(in_features=input_output, out_features=input_output)
 
     def forward(self, x):
         """Forward pass."""
         out = self.fc1(x)
         out = self.sigmoid1(out)
         out = self.fc2(out)
-        out = self.sigmoid2(out)
-        out = self.fc3(out)
-        out = self.sigmoid3(out)
-        out = self.fc4(out)
-        out = self.sigmoid4(out)
-        out = self.fc5(out)
 
         return out
 
 
 @pytest.mark.parametrize(
-    "model, input_shape",
-    [
-        pytest.param(FC, (100, 32 * 32 * 3)),
-    ],
+    "model",
+    [pytest.param(FC)],
 )
-def test_compile_torch(model, input_shape, default_compilation_configuration, seed_torch):
+@pytest.mark.parametrize(
+    "input_output_feature",
+    [pytest.param(input_output_feature) for input_output_feature in INPUT_OUTPUT_FEATURE],
+)
+def test_compile_torch(input_output_feature, model, seed_torch, default_compilation_configuration):
     """Test the different model architecture from torch numpy."""
 
     # Seed torch
     seed_torch()
 
-    # Define the torch model
-    torch_fc_model = model()
+    n_bits = 2
 
+    # Define an input shape (n_examples, n_features)
+    n_examples = 10
+
+    # Define the torch model
+    torch_fc_model = model(input_output_feature)
     # Create random input
-    torch_inputset = torch.randn(input_shape)
+    inputset = [numpy.random.uniform(-1, 1, size=input_output_feature) for _ in range(n_examples)]
 
     # Compile
-    compile_torch_model(
+    quantized_numpy_module = compile_torch_model(
         torch_fc_model,
-        torch_inputset,
+        inputset,
         default_compilation_configuration,
+        n_bits=n_bits,
     )
+
+    # Compare predictions between FHE and QuantizedModule
+    clear_predictions = []
+    homomorphic_predictions = []
+    for numpy_input in inputset:
+        q_input = QuantizedArray(n_bits, numpy_input)
+        x_q = q_input.qvalues
+        clear_predictions.append(quantized_numpy_module.forward(x_q))
+        homomorphic_predictions.append(
+            quantized_numpy_module.forward_fhe.run(numpy.array([x_q]).astype(numpy.uint8))
+        )
+
+    clear_predictions = numpy.array(clear_predictions)
+    homomorphic_predictions = numpy.array(homomorphic_predictions)
+
+    # Make sure homomorphic_predictions are the same as dequant_predictions
+    assert numpy.array_equal(homomorphic_predictions, clear_predictions)