From 25e12530d18d5348c592da0af41cccada25b5687 Mon Sep 17 00:00:00 2001 From: jfrery Date: Thu, 16 Dec 2021 18:24:27 +0100 Subject: [PATCH] feat: make the last layer over 7 bits --- concrete/quantization/post_training.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/concrete/quantization/post_training.py b/concrete/quantization/post_training.py index f60fb8614..4736fcb1b 100644 --- a/concrete/quantization/post_training.py +++ b/concrete/quantization/post_training.py @@ -88,7 +88,12 @@ class PostTrainingAffineQuantization: # Create a QuantizedLinear layer q_weights = self.quant_params[f"{name}.weight"] q_bias = self.quant_params[f"{name}.bias"] - q_layer = QuantizedLinear(self.n_bits, q_weights, q_bias) + # Check if layer is last layer from the model + if name == list(self.numpy_model.torch_model.named_children())[-1][0]: + # If last layer, we can use 7 bits (maximum allowed) of precision. + q_layer = QuantizedLinear(7, q_weights, q_bias) + else: + q_layer = QuantizedLinear(self.n_bits, q_weights, q_bias) # Calibrate and get new calibration_data for next layer/activation calibration_data = self._calibrate_layers_activation( name, q_layer, calibration_data