diff --git a/python/treelite/model.py b/python/treelite/model.py index 9b4838b7..4e551143 100644 --- a/python/treelite/model.py +++ b/python/treelite/model.py @@ -3,6 +3,7 @@ from __future__ import annotations import ctypes +import json import pathlib import platform from typing import Any, List, Optional, Union @@ -72,6 +73,24 @@ def output_type(self) -> str: _check_call(_LIB.TreeliteGetOutputType(self.handle, ctypes.byref(out))) return py_str(out.value) + @property + def attributes(self) -> dict[Any, Any]: + """Optional model attributes (JSON string)""" + if self.handle is None: + raise AttributeError("Model not loaded yet") + + obj = _TreelitePyBufferFrame() + _check_call( + _LIB.TreeliteGetHeaderField( + self.handle, + c_str("attributes"), + ctypes.byref(obj), + ) + ) + array = _pybuffer2numpy(obj) + attributes_str = array.tobytes().decode("utf-8") + return json.loads(attributes_str) + @classmethod def concatenate(cls, model_objs: List[Model]) -> Model: """ diff --git a/python/treelite/model_builder.py b/python/treelite/model_builder.py index c4021d4a..2b042a30 100644 --- a/python/treelite/model_builder.py +++ b/python/treelite/model_builder.py @@ -361,7 +361,7 @@ def data_count(self, data_count: int): Number of data points """ _check_call( - _LIB.TreeliteModelBuilderGain( + _LIB.TreeliteModelBuilderDataCount( self.handle, ctypes.c_uint64(data_count), ) diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index b8c9cd30..ee0678a5 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -1,5 +1,6 @@ """Converter to export Treelite models as scikit-learn models (EXPERIMENTAL)""" +import warnings from enum import IntEnum from typing import Any @@ -103,8 +104,25 @@ def _export_tree( nodes["feature"] = tree_accessor.get_field("split_index") nodes["threshold"] = tree_accessor.get_field("threshold") nodes["impurity"] = np.nan - nodes["n_node_samples"] = -1 - nodes["weighted_n_node_samples"] = np.nan + data_count = tree_accessor.get_field("data_count").astype(np.intp) + data_count_mask = tree_accessor.get_field("data_count_present").astype(np.bool_) + if data_count.size == 0: + nodes["n_node_samples"] = np.full((n_nodes,), fill_value=-1, dtype=np.intp) + else: + data_count[~data_count_mask] = -1 + nodes["n_node_samples"] = data_count + # TODO(chyunsu3): In Treelite 5.0, rename field sum_hess -> weighted_data_count + weighted_data_count = tree_accessor.get_field("sum_hess").astype(np.float64) + weighted_data_count_mask = tree_accessor.get_field("sum_hess_present").astype( + np.bool_ + ) + if weighted_data_count.size == 0: + nodes["weighted_n_node_samples"] = np.full( + (n_nodes,), fill_value=np.nan, dtype=np.float64 + ) + else: + weighted_data_count[~weighted_data_count_mask] = np.nan + nodes["weighted_n_node_samples"] = weighted_data_count nodes["missing_go_to_left"] = tree_accessor.get_field("default_left") if n_targets == 1 and n_classes[0] == 1: @@ -154,7 +172,8 @@ def export_model(model: Model) -> Any: Note ---- - Currently only random forests can be exported as scikit-learn model objects. + Currently only random forests and isolation forests can be exported as + scikit-learn model objects. Support for gradient boosted trees and other kinds of tree models will be added in the future. @@ -168,15 +187,23 @@ def export_model(model: Model) -> Any: sklearn_model : object of type \ :py:class:`~sklearn.ensemble.RandomForestRegressor` / \ :py:class:`~sklearn.ensemble.RandomForestClassifier` / \ - :py:class:`~sklearn.ensemble.GradientBoostingRegressor` / \ - :py:class:`~sklearn.ensemble.GradientBoostingClassifier` + :py:class:`~sklearn.ensemble.IsolationForest` Scikit-learn model """ # pylint: disable=too-many-locals try: from sklearn import __version__ as sklearn_version - from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor - from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor + from sklearn.ensemble import ( + IsolationForest, + RandomForestClassifier, + RandomForestRegressor, + ) + from sklearn.ensemble._iforest import _average_path_length + from sklearn.tree import ( + DecisionTreeClassifier, + DecisionTreeRegressor, + ExtraTreeRegressor, + ) except ImportError as e: raise TreeliteError("This function requires scikit-learn package") from e @@ -225,6 +252,9 @@ def raise_not_rf_error(reason): if task_type in [_TaskType.kBinaryClf, _TaskType.kMultiClf]: estimator_class = RandomForestClassifier subestimator_class = DecisionTreeClassifier + elif task_type == _TaskType.kIsolationForest: + estimator_class = IsolationForest + subestimator_class = ExtraTreeRegressor else: estimator_class = RandomForestRegressor subestimator_class = DecisionTreeRegressor @@ -266,6 +296,42 @@ def raise_not_rf_error(reason): "classes_": [np.arange(n_classes[i]) for i in range(n_targets)], } ) + elif estimator_class is IsolationForest: + # Recover the `offset_` field; if missing, set to -0.5 + try: + offset = model.attributes["sklearn_iforest_offset"] + except KeyError: + warnings.warn( + "Treelite model does not store attribute 'sklearn_iforest_offset'; " + "setting it to the default value of -0.5...", + UserWarning, + ) + offset = -0.5 + + # Compute max_samples by taking the max over the weighted root counts + # (with bootstrap=True the unweighted root only counts distinct rows) + max_samples = int( + max(estimator.tree_.weighted_n_node_samples[0] for estimator in estimators) + ) + state.update( + { + "_max_samples": max_samples, + "max_samples_": max_samples, + "offset_": offset, + "_average_path_length_per_tree": tuple( + _average_path_length(est.tree_.n_node_samples) for est in estimators + ), + "_decision_path_lengths": tuple( + est.tree_.compute_node_depths() for est in estimators + ), + # The exported trees reference features globally, so scoring uses + # the full feature set for every tree. + "_max_features": n_features, + "estimators_features_": [ + np.arange(n_features, dtype=np.int64) for _ in estimators + ], + } + ) clf.__setstate__(state) return clf diff --git a/python/treelite/sklearn/importer.py b/python/treelite/sklearn/importer.py index 94b1a93b..8b2c0bc9 100644 --- a/python/treelite/sklearn/importer.py +++ b/python/treelite/sklearn/importer.py @@ -1,14 +1,15 @@ """Converter to ingest scikit-learn models into Treelite""" import ctypes +import json from typing import Optional import numpy as np from packaging.version import parse as parse_version from ..core import _LIB, TreeliteError, _check_call -from ..model import Model -from ..util import c_array +from ..model import Model, _numpy2pybuffer +from ..util import c_array, c_str from .isolation_forest import calculate_depths, expected_depth @@ -79,6 +80,18 @@ def import_model(sklearn_model) -> Model: # clf is an IsolationForest # tl_model is a Treelite representation of clf + To reproduce the output of :py:meth:`~sklearn.ensemble.IsolationForest.decision_function`, + retrieve the value of ``IsolationForest.offset_`` and apply it, + as follows: + + .. code-block:: python + + # Treelite model stores an optional list of attributes (as a JSON string). + # We can retrieve `offset_` from it. + offset = tl_model.attributes.get("sklearn_iforest_offset", -0.5) + # Apply offset_ to compute the decision function. + decision_function = -treelite.gtil.predict(tl_model, X) - offset + Parameters ---------- sklearn_model : object of type \ @@ -246,6 +259,7 @@ def import_model(sklearn_model) -> Model: ) ) elif isinstance(sklearn_model, IsolationForest): + # TODO(chyunsu3): In Treelite 5.0, pass offset_ field via TreeliteLoadSKLearnIsolationForest() _check_call( _LIB.TreeliteLoadSKLearnIsolationForest( ctypes.c_int(sklearn_model.n_estimators), @@ -263,6 +277,20 @@ def import_model(sklearn_model) -> Model: ctypes.byref(handle), ) ) + # Store `offset_` field as a model attribute + attributes = { + "sklearn_iforest_offset": float(sklearn_model.offset_), + } + attributes_serialized = json.dumps(attributes) + _check_call( + _LIB.TreeliteSetHeaderField( + handle, + c_str("attributes"), + _numpy2pybuffer( + np.frombuffer(attributes_serialized.encode("utf-8"), dtype="S1") + ), + ) + ) elif isinstance(sklearn_model, (RandomForestC, ExtraTreesC)): n_classes = np.array(sklearn_model.n_classes_, dtype=np.int32) _check_call( diff --git a/tests/python/test_model_builder.py b/tests/python/test_model_builder.py index 39937a8e..4183af05 100644 --- a/tests/python/test_model_builder.py +++ b/tests/python/test_model_builder.py @@ -139,3 +139,48 @@ def make_tree_stump(left_child_val, right_child_val): expected_pred = np.array([[2, 2], [1, 1]]) pred = treelite.gtil.predict_leaf(model, dmat) np.testing.assert_almost_equal(pred, expected_pred, decimal=5) + + +def test_data_count_setter(): + """Test whether data count can be specified as part of model builder""" + # Tree stump with 3 nodes + builder = ModelBuilder( + threshold_type="float32", + leaf_output_type="float32", + metadata=Metadata( + num_feature=2, + task_type="kRegressor", + average_tree_output=False, + num_target=1, + num_class=[1], + leaf_vector_shape=(1, 1), + ), + tree_annotation=TreeAnnotation(num_tree=1, target_id=[0], class_id=[0]), + postprocessor=PostProcessorFunc(name="identity"), + base_scores=[0.0], + ) + builder.start_tree() + builder.start_node(0) + builder.numerical_test( + feature_id=0, + threshold=0.0, + default_left=False, + opname="<=", + left_child_key=1, + right_child_key=2, + ) + builder.data_count(100) + builder.end_node() + builder.start_node(1) + builder.leaf(-1.0) + builder.data_count(10) + builder.end_node() + builder.start_node(2) + builder.leaf(1.0) + builder.data_count(90) + builder.end_node() + builder.end_tree() + + model = builder.commit() + data_count = model.get_tree_accessor(0).get_field("data_count") + np.testing.assert_array_equal(data_count, np.array([100, 10, 90], dtype=np.int32)) diff --git a/tests/python/test_sklearn_integration.py b/tests/python/test_sklearn_integration.py index 3e53d7d7..35067472 100644 --- a/tests/python/test_sklearn_integration.py +++ b/tests/python/test_sklearn_integration.py @@ -187,11 +187,74 @@ def test_skl_converter_iforest(dataset): random_state=0, ) clf.fit(X) - expected_pred = -clf.score_samples(X) - expected_pred = expected_pred.reshape((-1, 1, 1)) + tl_model = treelite.sklearn.import_model(clf) + + # 1. Compare raw anomaly scores + np.testing.assert_almost_equal( + -treelite.gtil.predict(tl_model, X), + clf.score_samples(X).reshape((-1, 1, 1)), + ) + + # 2. Compare decision_function + # (decision_function = score_samples - offset) + offset = tl_model.attributes["sklearn_iforest_offset"] + np.testing.assert_almost_equal( + -treelite.gtil.predict(tl_model, X) - offset, + clf.decision_function(X).reshape((-1, 1, 1)), + ) + +@pytest.mark.parametrize("bootstrap", [True, False]) +@pytest.mark.parametrize("use_sample_weights", [True, False]) +def test_iforest_round_trip(bootstrap, use_sample_weights): + """ + Ensure that Treelite preserve important attributes when importing + and exporting isolation forests. + """ + + n_samples, n_outliers = 120, 40 + rng = np.random.RandomState(0) + covariance = np.array([[0.5, -0.1], [0.7, 0.4]]) + cluster_1 = 0.4 * rng.randn(n_samples, 2) @ covariance + np.array([2, 2]) + cluster_2 = 0.3 * rng.randn(n_samples, 2) + np.array([-2, -2]) + outliers = rng.uniform(low=-4, high=4, size=(n_outliers, 2)) + + X = np.concatenate([cluster_1, cluster_2, outliers]) + + clf = IsolationForest( + max_samples=100, + n_estimators=100, + n_jobs=-1, + random_state=0, + bootstrap=bootstrap, + ) + if use_sample_weights: + clf.fit(X, sample_weight=rng.uniform(low=0.2, high=0.8, size=(X.shape[0],))) + else: + clf.fit(X) tl_model = treelite.sklearn.import_model(clf) - out_pred = treelite.gtil.predict(tl_model, X) + exported_model = treelite.sklearn.export_model(tl_model) + assert type(exported_model) is type(clf) + assert len(clf.estimators_) == len(exported_model.estimators_) + for old_tree, new_tree in zip(clf.estimators_, exported_model.estimators_): + assert type(old_tree) is type(new_tree) + np.testing.assert_array_equal( + old_tree.tree_.n_node_samples, new_tree.tree_.n_node_samples + ) + np.testing.assert_almost_equal( + old_tree.tree_.weighted_n_node_samples, + new_tree.tree_.weighted_n_node_samples, + decimal=5, + ) + np.testing.assert_almost_equal(clf.offset_, exported_model.offset_) + np.testing.assert_almost_equal(clf.max_samples_, exported_model.max_samples_) + + expected_pred = clf.score_samples(X) + out_pred = exported_model.score_samples(X) + np.testing.assert_almost_equal(out_pred, expected_pred) + + expected_pred = clf.decision_function(X) + out_pred = exported_model.decision_function(X) np.testing.assert_almost_equal(out_pred, expected_pred)