Support count:poisson for XGBRegressor (#666)

Signed-off-by: Xavier Dupre <[email protected]>
onnx · Nov 23, 2023 · 1984712 · 1984712
1 parent 0b5bcf1
commit 1984712
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 7 deletions.
diff --git a/onnxmltools/convert/xgboost/operator_converters/XGBoost.py b/onnxmltools/convert/xgboost/operator_converters/XGBoost.py
@@ -250,19 +250,26 @@ def convert(scope, operator, container):
         )
 
         # add nodes
+        if objective == "count:poisson":
+            names = [scope.get_unique_variable_name("tree")]
+            del attr_pairs["base_values"]
+        else:
+            names = operator.output_full_names
         container.add_node(
             "TreeEnsembleRegressor",
             operator.input_full_names,
-            operator.output_full_names,
+            names,
             op_domain="ai.onnx.ml",
             name=scope.get_unique_operator_name("TreeEnsembleRegressor"),
             **attr_pairs,
         )
-        # try:
-        #    if len(inputs[0].type.tensor_type.shape.dim) > 0:
-        #        output_dim = [inputs[0].type.tensor_type.shape.dim[0].dim_value, 1]
-        # except Exception:
-        #    raise ValueError('Invalid/missing input dimension.')
+
+        if objective == "count:poisson":
+            cst = scope.get_unique_variable_name("half")
+            container.add_initializer(cst, TensorProto.FLOAT, [1], [0.5])
+            new_name = scope.get_unique_variable_name("exp")
+            container.add_node("Exp", names, [new_name])
+            container.add_node("Mul", [new_name, cst], operator.output_full_names)
 
 
 class XGBClassifierConverter(XGBConverter):

diff --git a/tests/xgboost/test_xgboost_converters.py b/tests/xgboost/test_xgboost_converters.py
@@ -89,6 +89,33 @@ def test_xgb_regressor(self):
             basename="SklearnXGBRegressor-Dec3",
         )
 
+    def test_xgb_regressor_poisson(self):
+        iris = load_diabetes()
+        x = iris.data
+        y = iris.target
+        x_train, x_test, y_train, _ = train_test_split(
+            x, y, test_size=0.5, random_state=42
+        )
+        for nest in [5, 50]:
+            xgb = XGBRegressor(
+                objective="count:poisson",
+                random_state=0,
+                max_depth=3,
+                n_estimators=nest,
+            )
+            xgb.fit(x_train, y_train)
+            conv_model = convert_xgboost(
+                xgb,
+                initial_types=[("input", FloatTensorType(shape=[None, None]))],
+                target_opset=TARGET_OPSET,
+            )
+            dump_data_and_model(
+                x_test.astype("float32"),
+                xgb,
+                conv_model,
+                basename=f"SklearnXGBRegressorPoisson{nest}-Dec3",
+            )
+
     def test_xgb_classifier(self):
         xgb, x_test = _fit_classification_model(XGBClassifier(), 2)
         conv_model = convert_xgboost(
@@ -650,5 +677,5 @@ def test_xgb_classifier_hinge(self):
 
 
 if __name__ == "__main__":
-    # TestXGBoostModels().test_xgboost_booster_classifier_multiclass_softprob()
+    TestXGBoostModels().test_xgb_regressor_poisson()
     unittest.main(verbosity=2)