In addition to Concrete ML models and custom models in torch, it is also possible to directly compile ONNX models. This can be particularly appealing, notably to import models trained with Keras.
ONNX models can be compiled by directly importing models that are already quantized with Quantization Aware Training (QAT) or by performing Post-Training Quantization (PTQ) with Concrete ML.
The following example shows how to compile an ONNX model using PTQ. The model was initially trained using Keras before being exported to ONNX. The training code is not shown here.
This example uses Post-Training Quantization, i.e., the quantization is not performed during training. This model would not have good performance in FHE. Quantization Aware Training should be added by the model developer. Additionally, importing QAT ONNX models can be done as shown below.
from concrete.ml.torch.compile import compile_onnx_model
from concrete.fhe.compilation import Configuration
"""A fully-connected model."""
hidden_layer_size = 10
output_size = 5
self.dense1 = tensorflow.keras.layers.Dense(
self.dense2 = tensorflow.keras.layers.Dense(output_size, activation=tensorflow.nn.relu6)
self.flatten = tensorflow.keras.layers.Flatten()
def call(self, inputs):
x = self.flatten(inputs)
x = self.dense1(x)
x = self.dense2(x)
n_bits = 6
input_output_feature = 2
input_shape = (input_output_feature,)
num_inputs = 1
n_examples = 5000
# Define the Keras model
keras_model = FC()
keras_model.build((None,) + input_shape)
# Create random input
input_set = numpy.random.uniform(-100, 100, size=(n_examples, *input_shape))
# Convert to ONNX
tf2onnx.convert.from_keras(keras_model, opset=14, output_path="tmp.model.onnx")
onnx_model = onnx.load("tmp.model.onnx")
quantized_module = compile_onnx_model(
onnx_model, input_set, n_bits=2
# Create test data from the same distribution and quantize using
# learned quantization parameters during compilation
x_test = tuple(numpy.random.uniform(-100, 100, size=(1, *input_shape)) for _ in range(num_inputs))
y_clear = quantized_module.forward(*x_test, fhe="disable")
y_fhe = quantized_module.forward(*x_test, fhe="execute")
print("Execution in clear: ", y_clear)
print("Execution in FHE: ", y_fhe)
print("Equality: ", numpy.sum(y_clear == y_fhe), "over", numpy.size(y_fhe), "values")
While Keras was used in this example, it is not officially supported. Additional work is needed to test all of Keras's types of layers and models.
QAT models contain quantizers in the ONNX graph. These quantizers ensure that the inputs to the Linear/Dense and Conv layers are quantized. Since these QAT models have quantizers that are configured during training to a specific number of bits, the ONNX graph will need to be imported using the same settings:
n_bits_qat = 3 # number of bits for weights and activations during training
quantized_numpy_module = compile_onnx_model(
The following operators are supported for evaluation and conversion to an equivalent FHE circuit. Other operators were not implemented, either due to FHE constraints or because they are rarely used in PyTorch activations or scikit-learn models.