From 93aaec4da92db4f648b63d33627e81a86ff86e02 Mon Sep 17 00:00:00 2001 From: Roman Bredehoft Date: Tue, 9 Apr 2024 14:18:50 +0200 Subject: [PATCH] chore: also fix shape coherence with sklearn --- src/concrete/ml/sklearn/base.py | 54 ++++++++++++++++++++-------- tests/sklearn/test_sklearn_models.py | 34 +++++++++--------- 2 files changed, 56 insertions(+), 32 deletions(-) diff --git a/src/concrete/ml/sklearn/base.py b/src/concrete/ml/sklearn/base.py index 4ced084582..7aea99ae08 100644 --- a/src/concrete/ml/sklearn/base.py +++ b/src/concrete/ml/sklearn/base.py @@ -694,6 +694,8 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray: Returns: numpy.ndarray: The post-processed predictions. """ + assert isinstance(y_preds, numpy.ndarray), "Output predictions must be an array." + return y_preds @@ -805,8 +807,12 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray: # If the prediction array is 1D, transform the output into a 2D array [1-p, p], # with p the initial output probabilities - if y_preds.ndim == 1 or y_preds.shape[1] == 1: - y_preds = numpy.concatenate((1 - y_preds, y_preds), axis=1) + # This is similar to what is done in scikit-learn + if y_preds.ndim == 1: + return numpy.vstack([1 - y_preds, y_preds]).T + + elif y_preds.shape[1] == 1: + return numpy.concatenate((1 - y_preds, y_preds), axis=1) # Else, apply the softmax operator else: @@ -1387,8 +1393,13 @@ def quantize_input(self, X: numpy.ndarray) -> numpy.ndarray: def dequantize_output(self, q_y_preds: numpy.ndarray) -> numpy.ndarray: self.check_model_is_fitted() - q_y_preds = self.output_quantizers[0].dequant(q_y_preds) - return q_y_preds + y_preds = self.output_quantizers[0].dequant(q_y_preds) + + # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn + if y_preds.ndim == 2 and y_preds.shape[1] == 1: + return y_preds.ravel() + + return y_preds def _get_module_to_compile(self) -> Union[Compiler, QuantizedModule]: assert self._tree_inference is not None, self._is_not_fitted_error_message() @@ -1442,7 +1453,12 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray: if not self._fhe_ensembling: y_preds = numpy.sum(y_preds, axis=-1) - assert_true(y_preds.ndim == 2, "y_preds should be a 2D array") + assert isinstance(y_preds, numpy.ndarray), "Output predictions must be an array." + + # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn + if y_preds.ndim == 2 and y_preds.shape[1] == 1: + return y_preds.ravel() + return y_preds return super().post_processing(y_preds) @@ -1693,6 +1709,10 @@ def dequantize_output(self, q_y_preds: numpy.ndarray) -> numpy.ndarray: # De-quantize the output values y_preds = self.output_quantizers[0].dequant(q_y_preds) + # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn + if y_preds.ndim == 2 and y_preds.shape[1] == 1: + return y_preds.ravel() + return y_preds def _get_module_to_compile(self) -> Union[Compiler, QuantizedModule]: @@ -1775,25 +1795,26 @@ def decision_function( """ # Here, we want to use SklearnLinearModelMixin's `predict` method as confidence scores are # the dot product's output values, without any post-processing - y_preds = SklearnLinearModelMixin.predict(self, X, fhe=fhe) - return y_preds + y_scores = SklearnLinearModelMixin.predict(self, X, fhe=fhe) + + return y_scores def predict_proba(self, X: Data, fhe: Union[FheMode, str] = FheMode.DISABLE) -> numpy.ndarray: - y_logits = self.decision_function(X, fhe=fhe) - y_proba = self.post_processing(y_logits) + y_scores = self.decision_function(X, fhe=fhe) + y_proba = self.post_processing(y_scores) return y_proba - # In scikit-learn, the argmax is done on the logits directly, not the probabilities + # In scikit-learn, the argmax is done on the scores directly, not the probabilities def predict(self, X: Data, fhe: Union[FheMode, str] = FheMode.DISABLE) -> numpy.ndarray: # Compute the predicted scores - y_logits = self.decision_function(X, fhe=fhe) + y_scores = self.decision_function(X, fhe=fhe) # Retrieve the class with the highest score # If there is a single dimension, only compare the scores to 0 - if y_logits.ndim == 1 or y_logits.shape[1] == 1: - y_preds = (y_logits > 0).astype(int) + if y_scores.ndim == 1: + y_preds = (y_scores > 0).astype(int) else: - y_preds = numpy.argmax(y_logits, axis=1) + y_preds = numpy.argmax(y_scores, axis=1) return self.classes_[y_preds] @@ -2002,6 +2023,11 @@ def dequantize_output(self, q_y_preds: numpy.ndarray) -> numpy.ndarray: self.check_model_is_fitted() # We compute the sorted argmax in FHE, which are integers. # No need to de-quantize the output values + + # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn + if q_y_preds.ndim > 1 and q_y_preds.shape[1] == 1: + return q_y_preds.ravel() + return q_y_preds def _get_module_to_compile(self) -> Union[Compiler, QuantizedModule]: diff --git a/tests/sklearn/test_sklearn_models.py b/tests/sklearn/test_sklearn_models.py index 9266a4dfa4..46cbed8f03 100644 --- a/tests/sklearn/test_sklearn_models.py +++ b/tests/sklearn/test_sklearn_models.py @@ -204,13 +204,11 @@ def check_correctness_with_sklearn( y_scores_sklearn = sklearn_model.decision_function(x) y_scores_fhe = model.decision_function(x, fhe=fhe) - # Currently, for single target data sets, Concrete models' outputs have shape (n, 1) - # while scikit-learn models' outputs have shape (n, ) - # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4029 - # assert y_scores_sklearn.shape == y_scores_fhe.shape, ( - # "Method 'decision_function' outputs different shapes between scikit-learn and " - # f"Concrete ML in FHE (fhe={fhe})" - # ) + assert y_scores_sklearn.shape == y_scores_fhe.shape, ( + "Method 'decision_function' outputs different shapes between scikit-learn and " + f"Concrete ML in FHE (fhe={fhe})" + ) + check_r2_score(y_scores_sklearn, y_scores_fhe, acceptance_score=acceptance_r2score) # LinearSVC models from scikit-learn do not provide a 'predict_proba' method @@ -231,13 +229,10 @@ def check_correctness_with_sklearn( y_pred_sklearn = sklearn_model.predict(x) y_pred_fhe = model.predict(x, fhe=fhe) - # Currently, for single target data sets, Concrete models' outputs have shape (n, 1) while - # scikit-learn models' outputs have shape (n, ) - # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4029 - # assert y_pred_sklearn.shape == y_pred_fhe.shape, ( - # "Method 'predict' outputs different shapes between scikit-learn and " - # f"Concrete ML in FHE (fhe={fhe})" - # ) + assert y_pred_sklearn.shape == y_pred_fhe.shape, ( + "Method 'predict' outputs different shapes between scikit-learn and " + f"Concrete ML in FHE (fhe={fhe})" + ) # If the model is a classifier, check that accuracies are similar if is_classifier_or_partial_classifier(model): @@ -649,10 +644,14 @@ def check_separated_inference(model, fhe_circuit, x, check_float_array_equal): and get_model_name(model) != "KNeighborsClassifier" ): # For linear classifiers, the argmax is done on the scores directly, not the probabilities + # Also, it is handled differently if shape is (n,) instead of (n, 1) if is_model_class_in_a_list(model, _get_sklearn_linear_models()): - y_pred = numpy.argmax(y_scores, axis=-1) + if y_scores.ndim == 1: + y_pred = (y_scores > 0).astype(int) + else: + y_pred = numpy.argmax(y_scores, axis=1) else: - y_pred = numpy.argmax(y_pred, axis=-1) + y_pred = numpy.argmax(y_pred, axis=1) y_pred_class = model.predict(x, fhe="simulate") @@ -1874,7 +1873,6 @@ def test_rounding_consistency_for_regular_models( n_bits, load_data, check_r2_score, - check_accuracy, is_weekly_option, verbose=True, ): @@ -1894,7 +1892,7 @@ def test_rounding_consistency_for_regular_models( else: # Check `predict` for regressors predict_method = model.predict - metric = check_accuracy + metric = check_r2_score check_rounding_consistency( model,