support objective=gamma in lgbm regressor

Signed-off-by: Jan-Benedikt Jagusch <jan.jagusch@gmail.com>
onnx · Jul 27, 2021 · d0d9f8b · d0d9f8b
1 parent f295db6
commit d0d9f8b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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,