From 19847127edd283dc0140e3fc5ec23e25a1f6dca6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Thu, 23 Nov 2023 01:20:54 +0100 Subject: [PATCH] Support count:poisson for XGBRegressor (#666) Signed-off-by: Xavier Dupre --- .../xgboost/operator_converters/XGBoost.py | 19 ++++++++---- tests/xgboost/test_xgboost_converters.py | 29 ++++++++++++++++++- 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/onnxmltools/convert/xgboost/operator_converters/XGBoost.py b/onnxmltools/convert/xgboost/operator_converters/XGBoost.py index 768c66ca..42a3964a 100644 --- 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 index 600456c0..4b1626ea 100644 --- 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)