Skip to content

Commit f089575

Browse files
authored
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 <[email protected]>
1 parent 2351850 commit f089575

File tree

2 files changed

+2
-1
lines changed

2 files changed

+2
-1
lines changed

onnxmltools/convert/lightgbm/operator_converters/LightGbm.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ def convert_lightgbm(scope, operator, container):
444444
elif gbm_text['objective'].startswith('multiclass'):
445445
n_classes = gbm_text['num_class']
446446
attrs['post_transform'] = 'SOFTMAX'
447-
elif gbm_text['objective'].startswith('regression'):
447+
elif gbm_text['objective'].startswith(('regression', 'quantile')):
448448
n_classes = 1 # Regressor has only one output variable
449449
attrs['post_transform'] = 'NONE'
450450
attrs['n_targets'] = n_classes

tests/lightgbm/test_objective_functions.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ class ObjectiveTest(unittest.TestCase):
3131
"regression",
3232
"poisson",
3333
"gamma",
34+
"quantile"
3435
)
3536

3637
@staticmethod

0 commit comments

Comments
 (0)