feat: end-to-end compilation of a torch model

tests/quantization/test_compilation.py

@@ -0,0 +1,84 @@
+"""Test Neural Networks compilations"""
+import numpy
+import pytest
+from torch import nn
+
+from concrete.quantization import PostTrainingAffineQuantization, QuantizedArray
+from concrete.torch import NumpyModule
+
+# 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, input_output):
+        super().__init__()
+        self.fc1 = nn.Linear(in_features=input_output, out_features=input_output)
+        self.sigmoid1 = nn.Sigmoid()
+        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)
+
+        return out
+
+
+@pytest.mark.parametrize(
+    "model",
+    [pytest.param(FC)],
+)
+@pytest.mark.parametrize(
+    "input_output_feature",
+    [pytest.param(input_output_feature) for input_output_feature in INPUT_OUTPUT_FEATURE],
+)
+def test_quantized_module_compilation(
+    input_output_feature, model, seed_torch, default_compilation_configuration
+):
+    """Test a neural network compilation for FHE inference."""
+    # Seed torch
+    seed_torch()
+
+    n_bits = 2
+
+    # Define an input shape (n_examples, n_features)
+    input_shape = (10, input_output_feature)
+
+    # Build a random Quantized Fully Connected Neural Network
+
+    # Define the torch model
+    torch_fc_model = model(input_output_feature)
+    # Create random input
+    numpy_input = numpy.random.uniform(-1, 1, size=input_shape)
+    # Create corresponding numpy model
+    numpy_fc_model = NumpyModule(torch_fc_model)
+    # Quantize with post-training static method
+    post_training_quant = PostTrainingAffineQuantization(n_bits, numpy_fc_model)
+    quantized_model = post_training_quant.quantize_module(numpy_input)
+    # Quantize input
+    q_input = QuantizedArray(n_bits, numpy_input)
+    quantized_model(q_input)
+
+    # Compile
+    quantized_model.compile(q_input, default_compilation_configuration)
+    dequant_predictions = quantized_model.forward_and_dequant(q_input)
+
+    # Compare predictions between FHE and QuantizedModule
+    homomorphic_predictions = []
+    for x_q in q_input.qvalues:
+        homomorphic_predictions.append(
+            quantized_model.forward_fhe.run(numpy.array([x_q]).astype(numpy.uint8))
+        )
+    homomorphic_predictions = quantized_model.dequantize_output(
+        numpy.array(homomorphic_predictions, dtype=numpy.float32)
+    )
+
+    homomorphic_predictions.reshape(dequant_predictions.shape)
+
+    # Make sure homomorphic_predictions are the same as dequant_predictions
+    assert numpy.isclose(homomorphic_predictions.ravel(), dequant_predictions.ravel()).all()

tests/quantization/test_quantized_activations.py

@@ -10,7 +10,7 @@ N_BITS_ATOL_TUPLE_LIST = [
     (20, 10 ** -2),
     (16, 10 ** -1),
     (8, 10 ** -0),
-    (4, 10 ** -0),
+    (5, 10 ** -0),
 ]
 
 

tests/quantization/test_quantized_layers.py

@@ -16,7 +16,7 @@ N_BITS_LIST = [20, 16, 8]
 @pytest.mark.parametrize(
     "n_examples, n_features, n_neurons",
     [
-        pytest.param(2, 3, 4),
+        pytest.param(50, 3, 4),
         pytest.param(20, 500, 30),
         pytest.param(200, 300, 50),
         pytest.param(10000, 100, 1),
@@ -33,7 +33,7 @@ def test_quantized_linear(n_examples, n_features, n_neurons, n_bits, is_signed):
     inputs = numpy.random.uniform(size=(n_examples, n_features))
     q_inputs = QuantizedArray(n_bits, inputs)
 
-    # shape of weights: (n_neurons, n_features)
+    # shape of weights: (n_features, n_neurons)
     weights = numpy.random.uniform(size=(n_features, n_neurons))
     q_weights = QuantizedArray(n_bits, weights, is_signed)
 
@@ -49,7 +49,7 @@ def test_quantized_linear(n_examples, n_features, n_neurons, n_bits, is_signed):
     expected_outputs = q_linear.q_out.values
     actual_output = q_linear(q_inputs).dequant()
 
-    assert numpy.isclose(expected_outputs, actual_output, rtol=10 ** -1).all()
+    assert numpy.isclose(expected_outputs, actual_output, atol=10 ** -0).all()
 
     # Same test without bias
     q_linear = QuantizedLinear(n_bits, q_weights)
@@ -59,4 +59,4 @@ def test_quantized_linear(n_examples, n_features, n_neurons, n_bits, is_signed):
     expected_outputs = q_linear.q_out.values
     actual_output = q_linear(q_inputs).dequant()
 
-    assert numpy.isclose(expected_outputs, actual_output, rtol=10 ** -1).all()
+    assert numpy.isclose(expected_outputs, actual_output, atol=10 ** -0).all()

tests/quantization/test_quantized_module.py

@@ -101,9 +101,7 @@ def test_quantized_linear(model, input_shape, n_bits, atol, seed_torch):
     quantized_model = post_training_quant.quantize_module(numpy_input)
     # Quantize input
     q_input = QuantizedArray(n_bits, numpy_input)
-    # Get quantized prediction
-    q_prediction = quantized_model(q_input)
-    # Dequantize to get back to real values
-    dequant_prediction = q_prediction.dequant()
+    # Forward and Dequantize to get back to real values
+    dequant_prediction = quantized_model.forward_and_dequant(q_input)
 
     assert numpy.isclose(numpy_prediction, dequant_prediction, atol=atol).all()