add support for quantile objective for LGBM models (#503)

* add support for quantile objective for LGBM models * add test case Signed-off-by: Mohammed Ajil <mohammed@ajil.ch>
onnx · Oct 4, 2021 · f089575 · f089575
1 parent 2351850
commit f089575
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py
@@ -444,7 +444,7 @@ def convert_lightgbm(scope, operator, container):
     elif gbm_text['objective'].startswith('multiclass'):
         n_classes = gbm_text['num_class']
         attrs['post_transform'] = 'SOFTMAX'
-    elif gbm_text['objective'].startswith('regression'):
+    elif gbm_text['objective'].startswith(('regression', 'quantile')):
         n_classes = 1  # Regressor has only one output variable
         attrs['post_transform'] = 'NONE'
         attrs['n_targets'] = n_classes

diff --git a/tests/lightgbm/test_objective_functions.py b/tests/lightgbm/test_objective_functions.py
@@ -31,6 +31,7 @@ class ObjectiveTest(unittest.TestCase):
         "regression",
         "poisson",
         "gamma",
+        "quantile"
     )
 
     @staticmethod