From d0d9f8b67da5b2f3555cd40f09a81b5d5a3f564d Mon Sep 17 00:00:00 2001 From: Jan-Benedikt Jagusch Date: Tue, 27 Jul 2021 17:44:38 +0200 Subject: [PATCH] support objective=gamma in lgbm regressor Signed-off-by: Jan-Benedikt Jagusch --- onnxmltools/convert/lightgbm/_parse.py | 3 ++- onnxmltools/convert/lightgbm/operator_converters/LightGbm.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py index 6c384e12..e527faf7 100644 --- a/onnxmltools/convert/lightgbm/_parse.py +++ b/onnxmltools/convert/lightgbm/_parse.py @@ -28,7 +28,8 @@ def __init__(self, booster): _model_dict['objective'].startswith('multiclass')): self.operator_name = 'LgbmClassifier' elif (_model_dict['objective'].startswith('regression') or - _model_dict['objective'].startswith('poisson')): + _model_dict['objective'].startswith('poisson') or + _model_dict['objective'].startswith('gamma')): self.operator_name = 'LgbmRegressor' else: # Other objectives are not supported. diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py index a7748799..c2ad360c 100644 --- a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py +++ b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py @@ -234,7 +234,7 @@ def convert_lightgbm(scope, operator, container): n_classes = 1 # Regressor has only one output variable attrs['post_transform'] = 'NONE' attrs['n_targets'] = n_classes - elif gbm_text['objective'].startswith('poisson'): + elif gbm_text['objective'].startswith('poisson') or gbm_text['objective'].startswith('gamma'): n_classes = 1 # Regressor has only one output variable attrs['n_targets'] = n_classes # 'Exp' is not a supported post_transform value in the ONNX spec yet,