Support xgboost 3

Signed-off-by: Asher Wright <asherw@squareup.com>

Author: AsherWright

onnxmltools/convert/xgboost/_parse.py

@@ -121,9 +121,16 @@ def _get_attributes(booster):
             kwargs["objective"] = "binary:logistic"
 
     if "base_score" not in kwargs:
-        kwargs["base_score"] = 0.5
+        kwargs["base_score"] = [0.5]
     elif isinstance(kwargs["base_score"], str):
-        kwargs["base_score"] = float(kwargs["base_score"])
+        base_score_str = kwargs["base_score"]
+        if base_score_str.startswith("[") and base_score_str.endswith("]"):
+            # xgboost >= 3.0: base_score is a string array
+            bs = json.loads(base_score_str)
+            kwargs["base_score"] = [float(x) for x in bs]
+        else:
+            # xgboost >= 2, < 3: base_score is a string float
+            kwargs["base_score"] = [float(base_score_str)]
     return kwargs
 
 

onnxmltools/convert/xgboost/common.py

@@ -29,9 +29,15 @@ def get_xgb_params(xgb_node):
         if xgb_node.n_estimators is not None:
             params["n_estimators"] = xgb_node.n_estimators
     if "base_score" in config["learner"]["learner_model_param"]:
-        bs = float(config["learner"]["learner_model_param"]["base_score"])
-        # xgboost >= 2.0
-        params["base_score"] = bs
+        base_score_raw = config["learner"]["learner_model_param"]["base_score"]
+        # xgboost >= 3.0: base_score is a string array
+        if base_score_raw.startswith("[") and base_score_raw.endswith("]"):
+            base_score = json.loads(base_score_raw)
+            params["base_score"] = [float(x) for x in base_score]
+        else:
+            # xgboost >= 2, < 3: base_score is a string float
+            params["base_score"] = [float(base_score_raw)]
+
     if "num_target" in config["learner"]["learner_model_param"]:
         params["n_targets"] = int(
             config["learner"]["learner_model_param"]["num_target"]

onnxmltools/convert/xgboost/operator_converters/XGBoost.py

@@ -53,7 +53,7 @@ def common_members(xgb_node, inputs):
         else:
             best_ntree_limit = params.get("best_ntree_limit", None)
         if base_score is None:
-            base_score = 0.5
+            base_score = [0.5]
         booster = xgb_node.get_booster()
         # The json format was available in October 2017.
         # XGBoost 0.7 was the first version released with it.
@@ -259,7 +259,10 @@ def convert(scope, operator, container):
             raise RuntimeError("Objective '{}' not supported.".format(objective))
 
         attr_pairs = XGBRegressorConverter._get_default_tree_attribute_pairs()
-        attr_pairs["base_values"] = [base_score]
+        if isinstance(base_score, list):
+            attr_pairs["base_values"] = base_score
+        else:
+            attr_pairs["base_values"] = [base_score]
 
         if best_ntree_limit and best_ntree_limit < len(js_trees):
             js_trees = js_trees[:best_ntree_limit]
@@ -288,7 +291,8 @@ def convert(scope, operator, container):
 
         if objective == "count:poisson":
             cst = scope.get_unique_variable_name("poisson")
-            container.add_initializer(cst, TensorProto.FLOAT, [1], [base_score])
+            # base_score is always a list
+            container.add_initializer(cst, TensorProto.FLOAT, [len(base_score)], base_score)
             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)
@@ -350,17 +354,27 @@ def convert(scope, operator, container):
                 attr_pairs["post_transform"] = "LOGISTIC"
                 attr_pairs["class_ids"] = [0 for v in attr_pairs["class_treeids"]]
                 if js_trees[0].get("leaf", None) == 0:
-                    attr_pairs["base_values"] = [base_score]
-                elif base_score != 0.5:
-                    # 0.5 -> cst = 0
-                    cst = -np.log(1 / np.float32(base_score) - 1.0)
-                    attr_pairs["base_values"] = [cst]
+                    # base_score is always a list
+                    attr_pairs["base_values"] = base_score
+                else:
+                    # Transform base_score - for binary, use first element
+                    bs_val = base_score[0]
+                    if bs_val != 0.5:
+                        # 0.5 -> cst = 0
+                        cst = -np.log(1 / np.float32(bs_val) - 1.0)
+                        attr_pairs["base_values"] = [cst]
             else:
-                attr_pairs["base_values"] = [base_score]
+                # base_score is always a list
+                attr_pairs["base_values"] = base_score
         else:
             # See https://github.com/dmlc/xgboost/blob/main/src/common/math.h#L35.
             attr_pairs["post_transform"] = "SOFTMAX"
-            attr_pairs["base_values"] = [base_score for n in range(ncl)]
+            # base_score is always a list - may have different values per class in xgboost 3+
+            # If base_score has fewer elements than classes, replicate to match
+            if len(base_score) == 1:
+                attr_pairs["base_values"] = base_score * ncl
+            else:
+                attr_pairs["base_values"] = base_score
             attr_pairs["class_ids"] = [v % ncl for v in attr_pairs["class_treeids"]]
 
         classes = xgb_node.classes_

tests/xgboost/test_xgboost_converters_base_score.py

@@ -4,7 +4,8 @@
 import numpy as np
 import scipy
 from numpy.testing import assert_almost_equal
-from sklearn.datasets import make_regression
+from sklearn.datasets import make_regression, make_classification
+import json
 
 try:
     from xgboost import XGBClassifier, XGBRegressor
@@ -137,6 +138,111 @@ def test_xgbclassifier_sparse_no_base_score(self):
         got = sess.run(None, feeds)[1]
         assert_almost_equal(expected.reshape((-1, 2)), got, decimal=4)
 
+    @unittest.skipIf(XGBRegressor is None, "xgboost is not available")
+    def test_xgbclassifier_multiclass_base_score(self):
+        """Test multiclass classifier - xgboost 3 can have different base_scores per class"""
+        X, y = make_classification(
+            n_samples=200, n_features=10, n_classes=3,
+            n_informative=5, n_redundant=0, random_state=42
+        )
+        X = X.astype(np.float32)
+
+        clf = XGBClassifier(n_estimators=3, max_depth=4, random_state=42)
+        clf.fit(X, y)
+        expected = clf.predict_proba(X).astype(np.float32)
+
+        onx = convert_xgboost(
+            clf,
+            initial_types=[("X", FloatTensorType(shape=[None, None]))],
+            target_opset=TARGET_OPSET,
+        )
+        feeds = {"X": X}
+
+        sess = InferenceSession(
+            onx.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        got = sess.run(None, feeds)[1]
+        assert_almost_equal(expected, got, decimal=4)
+
+    @unittest.skipIf(XGBRegressor is None, "xgboost is not available")
+    def test_xgbclassifier_multiclass_base_score_in_onnx(self):
+        """Verify that base_values are actually present in the ONNX graph"""
+        X, y = make_classification(
+            n_samples=200, n_features=10, n_classes=3,
+            n_informative=5, n_redundant=0, random_state=42
+        )
+        X = X.astype(np.float32)
+
+        clf = XGBClassifier(n_estimators=3, max_depth=4, random_state=42)
+        clf.fit(X, y)
+
+        config = json.loads(clf.get_booster().save_config())
+        base_score_str = config["learner"]["learner_model_param"]["base_score"]
+
+        onx = convert_xgboost(
+            clf,
+            initial_types=[("X", FloatTensorType(shape=[None, None]))],
+            target_opset=TARGET_OPSET,
+        )
+
+        tree_ensemble_node = None
+        for node in onx.graph.node:
+            if node.op_type == "TreeEnsembleClassifier":
+                tree_ensemble_node = node
+                break
+
+        self.assertIsNotNone(tree_ensemble_node, "TreeEnsembleClassifier node not found")
+
+        base_values = None
+        for attr in tree_ensemble_node.attribute:
+            if attr.name == "base_values":
+                base_values = list(attr.floats)
+                break
+
+        self.assertIsNotNone(base_values, "base_values attribute not found in ONNX model")
+        self.assertEqual(len(base_values), 3, "base_values should have 3 elements for 3-class problem")
+
+        # In xgboost 3+, base_score is a string array like "[3.4E-1,3.3E-1,3.3E-1]"
+        # Verify that base_values in ONNX match the xgboost config
+        if base_score_str.startswith("[") and base_score_str.endswith("]"):
+            expected_base_scores = json.loads(base_score_str)
+            for i, val in enumerate(base_values):
+                if i < len(expected_base_scores):
+                    self.assertAlmostEqual(val, expected_base_scores[i], places=5)
+
+    @unittest.skipIf(XGBRegressor is None, "xgboost is not available")
+    def test_xgbregressor_base_score_in_onnx(self):
+        """Verify that regressor base_values are present in the ONNX graph"""
+        X, y = make_regression(n_samples=200, n_features=10, random_state=42)
+        X = X.astype(np.float32)
+        y = y.astype(np.float32)
+
+        reg = XGBRegressor(n_estimators=3, max_depth=4, random_state=42)
+        reg.fit(X, y)
+
+        onx = convert_xgboost(
+            reg,
+            initial_types=[("X", FloatTensorType(shape=[None, None]))],
+            target_opset=TARGET_OPSET,
+        )
+
+        tree_ensemble_node = None
+        for node in onx.graph.node:
+            if node.op_type == "TreeEnsembleRegressor":
+                tree_ensemble_node = node
+                break
+
+        self.assertIsNotNone(tree_ensemble_node, "TreeEnsembleRegressor node not found")
+
+        base_values = None
+        for attr in tree_ensemble_node.attribute:
+            if attr.name == "base_values":
+                base_values = list(attr.floats)
+                break
+
+        self.assertIsNotNone(base_values, "base_values attribute not found in ONNX model")
+        self.assertGreater(len(base_values), 0, "base_values should not be empty")
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)