From 324821652ec5cd8e2ba67478843da5afa12f09f3 Mon Sep 17 00:00:00 2001 From: kcelia Date: Tue, 23 Jan 2024 18:20:45 +0100 Subject: [PATCH] chore: update v3 --- src/concrete/ml/sklearn/tree_to_numpy.py | 16 ++++++++-------- tests/sklearn/test_sklearn_models.py | 4 ++-- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/concrete/ml/sklearn/tree_to_numpy.py b/src/concrete/ml/sklearn/tree_to_numpy.py index d3da342fe..d2723e925 100644 --- a/src/concrete/ml/sklearn/tree_to_numpy.py +++ b/src/concrete/ml/sklearn/tree_to_numpy.py @@ -136,25 +136,25 @@ def assert_add_node_and_constant_in_xgboost_regressor_graph(onnx_model: onnx.Mod ) -def add_transpose_after_last_node(onnx_model: onnx.ModelProto, use_fhe_sum: bool): +def add_transpose_after_last_node(onnx_model: onnx.ModelProto, fhe_ensembling: bool): """Add transpose after last node. Args: onnx_model (onnx.ModelProto): The ONNX model. - use_fhe_sum (bool): Determines whether the sum of the trees' outputs is computed in FHE. + fhe_ensembling (bool): Determines whether the sum of the trees' outputs is computed in FHE. Default to False. """ # Get the output node output_node = onnx_model.graph.output[0] - # The state of the 'use_fhe_sum' variable affects the structure of the model's ONNX graph. + # The state of the 'fhe_ensembling' variable affects the structure of the model's ONNX graph. # When the option is enabled, the graph is cut after the ReduceSum node. # When it is disabled, the graph is cut at the ReduceSum node, which alters the output shape. # Therefore, it is necessary to adjust this shape with the correct permutation. # When using FHE sum for tree ensembles, create the node with perm attribute equal to (1, 0) # Otherwise, create the node with perm attribute equal to (2, 1, 0) - perm = [1, 0] if use_fhe_sum else [2, 1, 0] + perm = [1, 0] if fhe_ensembling else [2, 1, 0] transpose_node = onnx.helper.make_node( "Transpose", @@ -221,7 +221,7 @@ def tree_onnx_graph_preprocessing( onnx_model: onnx.ModelProto, framework: str, expected_number_of_outputs: int, - use_fhe_sum: bool = False, + fhe_ensembling: bool = False, ): """Apply pre-processing onto the ONNX graph. @@ -230,7 +230,7 @@ def tree_onnx_graph_preprocessing( framework (str): The framework from which the ONNX model is generated. (options: 'xgboost', 'sklearn') expected_number_of_outputs (int): The expected number of outputs in the ONNX model. - use_fhe_sum (bool): Determines whether the sum of the trees' outputs is computed in FHE. + fhe_ensembling (bool): Determines whether the sum of the trees' outputs is computed in FHE. Default to False. """ # Make sure the ONNX version returned by Hummingbird is OPSET_VERSION_FOR_ONNX_EXPORT @@ -258,7 +258,7 @@ def tree_onnx_graph_preprocessing( # Cut the graph after the ReduceSum node to remove # argmax, sigmoid, softmax from the graph. - if use_fhe_sum: + if fhe_ensembling: clean_graph_after_node_op_type(onnx_model, "ReduceSum") else: clean_graph_at_node_op_type(onnx_model, "ReduceSum") @@ -274,7 +274,7 @@ def tree_onnx_graph_preprocessing( # sklearn models apply the reduce sum before the transpose. # To have equivalent output between xgboost in sklearn, # apply the transpose before returning the output. - add_transpose_after_last_node(onnx_model, use_fhe_sum) + add_transpose_after_last_node(onnx_model, fhe_ensembling) # Cast nodes are not necessary so remove them. remove_node_types(onnx_model, op_types_to_remove=["Cast"]) diff --git a/tests/sklearn/test_sklearn_models.py b/tests/sklearn/test_sklearn_models.py index 446ad3d11..cd156f808 100644 --- a/tests/sklearn/test_sklearn_models.py +++ b/tests/sklearn/test_sklearn_models.py @@ -148,8 +148,8 @@ def get_n_bits_non_correctness(model_class): # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/3979 if get_model_name(model_class) == "KNeighborsClassifier": n_bits = 2 - - n_bits = min(N_BITS_REGULAR_BUILDS) + else: + n_bits = min(N_BITS_REGULAR_BUILDS) return n_bits