From 416aa913227afcc6ab42a6d618bec447f717c74e Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Fri, 21 Aug 2026 14:34:00 -0700 Subject: [PATCH 1/2] Exporter: export classes_ for subestimator of RF --- python/treelite/sklearn/exporter.py | 14 +++++++++-- tests/python/test_sklearn_integration.py | 32 +++++++++++++++++++++++- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index ee0678a5..62b6cd8a 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -159,9 +159,19 @@ def _export_tree( } if subestimator_class is DecisionTreeClassifier: if n_targets == 1: - subestimator_state["n_classes_"] = n_classes[0] + subestimator_state.update( + { + "n_classes_": n_classes[0], + "classes_": np.arange(n_classes[0]), + } + ) else: - subestimator_state["n_classes_"] = n_classes.tolist() + subestimator_state.update( + { + "n_classes_": n_classes.tolist(), + "classes_": [np.arange(n_classes[i]) for i in range(n_targets)], + } + ) subestimator.__setstate__(subestimator_state) return subestimator diff --git a/tests/python/test_sklearn_integration.py b/tests/python/test_sklearn_integration.py index 35067472..807827b3 100644 --- a/tests/python/test_sklearn_integration.py +++ b/tests/python/test_sklearn_integration.py @@ -236,7 +236,9 @@ def test_iforest_round_trip(bootstrap, use_sample_weights): 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_): + for old_tree, new_tree in zip( + clf.estimators_, exported_model.estimators_, strict=True + ): 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 @@ -417,3 +419,31 @@ def test_skl_export_rf_multitarget_multiclass(n_classes, n_estimators): clf2 = treelite.sklearn.export_model(tl_model) assert isinstance(clf2, RandomForestClassifier) np.testing.assert_almost_equal(clf2.predict_proba(X), clf.predict_proba(X)) + + +@given( + dataset=standard_classification_datasets( + n_classes=integers(min_value=2, max_value=4), + ), + n_estimators=integers(min_value=5, max_value=10), +) +@settings(**standard_settings()) +def test_random_forest_classifier_round_trip(dataset, n_estimators): + X, y = dataset + clf = RandomForestClassifier(random_state=0, n_estimators=n_estimators, n_jobs=-1) + clf.fit(X, y) + + tl_model = treelite.sklearn.import_model(clf) + clf2 = treelite.sklearn.export_model(tl_model) + + assert type(clf) is type(clf2) + assert len(clf.estimators_) == len(clf2.estimators_) + + np.testing.assert_array_equal(clf.classes_, clf2.classes_) + for old_tree, new_tree in zip(clf.estimators_, clf2.estimators_, strict=True): + assert type(old_tree) is type(new_tree) + assert old_tree.n_classes_ == new_tree.n_classes_ + np.testing.assert_array_equal( + old_tree.tree_.n_node_samples, new_tree.tree_.n_node_samples + ) + np.testing.assert_array_equal(old_tree.classes_, new_tree.classes_) From fd2187d7106a91eebd62c895bc0b8195ce15ee90 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 25 Aug 2026 17:03:38 -0700 Subject: [PATCH 2/2] Use np.float64 for classes_ --- python/treelite/sklearn/exporter.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index 62b6cd8a..6aaa2e51 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -162,14 +162,17 @@ def _export_tree( subestimator_state.update( { "n_classes_": n_classes[0], - "classes_": np.arange(n_classes[0]), + "classes_": np.arange(n_classes[0], dtype=np.float64), } ) else: subestimator_state.update( { "n_classes_": n_classes.tolist(), - "classes_": [np.arange(n_classes[i]) for i in range(n_targets)], + "classes_": [ + np.arange(n_classes[i], dtype=np.float64) + for i in range(n_targets) + ], } ) subestimator.__setstate__(subestimator_state)