8.5 KiB
module torus.ml.sklearn.tree_to_numpy
Implements the conversion of a tree model to a numpy function.
Global Variables
- MAX_BITWIDTH_BACKWARD_COMPATIBLE
- OPSET_VERSION_FOR_ONNX_EXPORT
- MSB_TO_KEEP_FOR_TREES
- MIN_CIRCUIT_THRESHOLD_FOR_TREES
function get_onnx_model
get_onnx_model(model, x: ndarray, framework: str) → ModelProto
Create ONNX model with Hummingbird convert method.
Args:
model: The tree model to convert.x(numpy.ndarray): Dataset used to trace the tree inference and convert the model to ONNX.framework(str): The framework from which the ONNX model is generated.(options: 'xgboost', 'sklearn')
Returns:
onnx.ModelProto: The ONNX model.
function workaround_squeeze_node_xgboost
workaround_squeeze_node_xgboost(onnx_model: ModelProto)
Workaround to fix torch issue that does not export the proper axis in the ONNX squeeze node.
FIXME: https://github.com/luxfi/torus-ml-internal/issues/2778 The squeeze ops does not have the proper dimensions. remove the following workaround when the issue is fixed Add the axis attribute to the Squeeze node
Args:
onnx_model(onnx.ModelProto): The ONNX model.
function assert_add_node_and_constant_in_xgboost_regressor_graph
assert_add_node_and_constant_in_xgboost_regressor_graph(onnx_model: ModelProto)
Assert if an Add node with a specific constant exists in the ONNX graph.
Args:
onnx_model(onnx.ModelProto): The ONNX model.
function add_transpose_after_last_node
add_transpose_after_last_node(
onnx_model: ModelProto,
fhe_ensembling: bool = False
)
Add transpose after last node.
Args:
onnx_model(onnx.ModelProto): The ONNX model.fhe_ensembling(bool): Determines whether the sum of the trees' outputs is computed in FHE. Default to False.
function preprocess_tree_predictions
preprocess_tree_predictions(
init_tensor: ndarray,
output_n_bits: int
) → QuantizedArray
Apply post-processing from the graph.
Args:
init_tensor(numpy.ndarray): Model parameters to be pre-processed.output_n_bits(int): The number of bits of the output.
Returns:
QuantizedArray: Quantizer for the tree predictions.
function tree_onnx_graph_preprocessing
tree_onnx_graph_preprocessing(
onnx_model: ModelProto,
framework: str,
expected_number_of_outputs: int,
fhe_ensembling: bool = False
)
Apply pre-processing onto the ONNX graph.
Args:
onnx_model(onnx.ModelProto): The ONNX model.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.fhe_ensembling(bool): Determines whether the sum of the trees' outputs is computed in FHE. Default to False.
function tree_values_preprocessing
tree_values_preprocessing(
onnx_model: ModelProto,
framework: str,
output_n_bits: int
) → QuantizedArray
Pre-process tree values.
Args:
onnx_model(onnx.ModelProto): The ONNX model.framework(str): The framework from which the ONNX model is generated.(options: 'xgboost', 'sklearn')output_n_bits(int): The number of bits of the output.
Returns:
QuantizedArray: Quantizer for the tree predictions.
function tree_to_numpy
tree_to_numpy(
model: BaseEstimator,
x: ndarray,
framework: str,
use_rounding: bool = True,
fhe_ensembling: bool = False,
output_n_bits: int = 8
) → Tuple[Callable, List[UniformQuantizer], ModelProto]
Convert the tree inference to a numpy functions using Hummingbird.
Args:
model(Callable): The tree model to convert.x(numpy.ndarray): The input data.use_rounding(bool): Determines whether the rounding feature is enabled or disabled. Default to True.fhe_ensembling(bool): Determines whether the sum of the trees' outputs is computed in FHE. Default to False.framework(str): The framework from which the ONNX model is generated.(options: 'xgboost', 'sklearn')output_n_bits(int): The number of bits of the output. Default to 8.
Returns:
Tuple[Callable, List[QuantizedArray], onnx.ModelProto]: A tuple with a function that takes a numpy array and returns a numpy array, QuantizedArray object to quantize and de-quantize the output of the tree, and the ONNX model.
function onnx_fp32_model_to_quantized_model
onnx_fp32_model_to_quantized_model(
onnx_model: ModelProto,
n_bits: int,
framework: str,
expected_number_of_outputs: int,
n_features: int,
model_inputs: Optional[ndarray] = None
)
Build a FHE-compliant onnx-model using a fitted scikit-learn model.
Args:
onnx_model(onnx.ModelProto): The fitted scikit-learn as a Hummingbird onnx model to convertn_bits(int): Number of bits to quantize the model. If an int is passed for n_bits, the value will be used for quantizing inputs and weights. If a dict is passed, then it should contain "op_inputs" and "op_weights" as keys with corresponding number of quantization bits so that: - op_inputs : number of bits to quantize the input values - op_weights: number of bits to quantize the learned parametersframework(str): either sklearn or xgboostexpected_number_of_outputs(int): expected number of outputsn_features(int): number of features as inputs of the modelmodel_inputs(Optional[numpy.ndarray]): optional dataset to use for quantization
Returns:
onnx.ModelProto: The converted onnx modelOptional[Tuple[int, int]]: Least significant bits to removelist[UniformQuantizer]: inputs quantizerslist[UniformQuantizer]: outputs quantizers