allow for identity POST_TRANSFORM (#674)

* allow for identity POST_TRANSFORM * add test
microsoft · Jan 10, 2023 · c862916 · c862916
1 parent a1fc5cf
commit c862916
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/hummingbird/ml/operator_converters/_gbdt_commons.py b/hummingbird/ml/operator_converters/_gbdt_commons.py
@@ -13,6 +13,7 @@
 from . import constants
 from ._tree_commons import get_tree_params_and_type, get_parameters_for_tree_trav_common, get_parameters_for_gemm_common
 from ._tree_commons import (
+    PostTransform,
     ApplySigmoidPostTransform,
     ApplySoftmaxPostTransform,
     ApplyTweediePostTransform,
@@ -160,6 +161,8 @@ def convert_gbdt_common(
                 extra_config[constants.POST_TRANSFORM] = ApplyTweedieBasePredictionPostTransform(base_prediction)
             else:
                 extra_config[constants.POST_TRANSFORM] = ApplyTweediePostTransform()
+        elif extra_config[constants.POST_TRANSFORM] is None:
+            extra_config[constants.POST_TRANSFORM] = PostTransform()
         else:
             raise NotImplementedError("Post transform {} not implemeneted yet".format(extra_config[constants.POST_TRANSFORM]))
     elif constants.BASE_PREDICTION in extra_config:

diff --git a/tests/test_lightgbm_converter.py b/tests/test_lightgbm_converter.py
@@ -88,6 +88,25 @@ def test_lgbm_tree_trav_multi_classifier_converter(self):
     def test_lgbm_perf_tree_trav_multi_classifier_converter(self):
         self._run_lgbm_classifier_converter(3, extra_config={"tree_implementation": "perf_tree_trav"})
 
+    @unittest.skipIf(not lightgbm_installed(), reason="LightGBM test requires LightGBM installed")
+    def test_lgbm_raw_logits(self):
+        warnings.filterwarnings("ignore")
+        model = lgb.LGBMClassifier(n_estimators=10)
+        np.random.seed(0)
+        X = np.random.rand(100, 200)
+        X = np.array(X, dtype=np.float32)
+        y = np.random.randint(3, size=100)
+
+        model.fit(X, y)
+
+        # "post_transform": None should return raw scores
+        torch_model = hummingbird.ml.convert(model, "torch", extra_config={"post_transform": None})
+        self.assertIsNotNone(torch_model)
+
+        np.testing.assert_allclose(
+            model.predict_proba(X, raw_score=True), torch_model.predict_proba(X), rtol=1e-06, atol=1e-06
+        )
+
     def _run_lgbm_ranker_converter(self, num_classes, extra_config={}, label_gain=None):
         warnings.filterwarnings("ignore")
         for max_depth in [1, 3, 8, 10, 12, None]: