Investigate a bug in xgboost (#520)

Infer num_class and n_estimators through tree_info
add target_opset to test cases
Adjust opset version for hummingbird due to torch opset version support
Signed-off-by: BowenBao <[email protected]>

Author: BowenBao

onnxmltools/convert/xgboost/_parse.py

@@ -27,11 +27,23 @@ def _append_covers(node):
 
 
 def _get_attributes(booster):
+    # num_class
+    state = booster.__getstate__()
+    bstate = bytes(state['handle'])
+    reg = re.compile(b'("tree_info":\\[[0-9,]*\\])')
+    objs = list(set(reg.findall(bstate)))
+    assert len(objs) == 1, 'Missing required property "tree_info".'
+    tree_info = json.loads("{{{}}}".format(objs[0].decode('ascii')))['tree_info']
+    num_class = len(set(tree_info))
+
     atts = booster.attributes()
-    ntrees = booster.best_ntree_limit
     dp = booster.get_dump(dump_format='json', with_stats=True)
     res = [json.loads(d) for d in dp]
     trees = len(res)
+    try:
+        ntrees = booster.best_ntree_limit
+    except AttributeError:
+        ntrees = trees // num_class if num_class > 0 else trees
     kwargs = atts.copy()
     kwargs['feature_names'] = booster.feature_names
     kwargs['n_estimators'] = ntrees
@@ -43,34 +55,22 @@ def _get_attributes(booster):
 
     if all(map(lambda x: int(x) == x, set(covs))):
         # regression
+        kwargs['num_target'] = num_class
         kwargs['num_class'] = 0
-        if trees > ntrees > 0:
-            kwargs['num_target'] = trees // ntrees
-            kwargs["objective"] = "reg:squarederror"
-        else:
-            kwargs['num_target'] = 1
-            kwargs["objective"] = "reg:squarederror"
+        kwargs["objective"] = "reg:squarederror"
     else:
         # classification
-        kwargs['num_target'] = 0
-        if trees > ntrees > 0:
-            state = booster.__getstate__()
-            bstate = bytes(state['handle'])
+        kwargs['num_class'] = num_class
+        if num_class != 1:
             reg = re.compile(b'(multi:[a-z]{1,15})')
             objs = list(set(reg.findall(bstate)))
-            if len(objs) != 1:
-                if '"name":"binary:logistic"' in str(bstate):
-                    kwargs['num_class'] = 1
-                    kwargs["objective"] = "binary:logistic"
-                else:
-                    raise RuntimeError(
-                        "Unable to guess objective in %r (trees=%r, ntrees=%r)"
-                        "." % (objs, trees, ntrees))
-            else:
-                kwargs['num_class'] = trees // ntrees
+            if len(objs) == 1:
                 kwargs["objective"] = objs[0].decode('ascii')
+            else:
+                raise RuntimeError(
+                    "Unable to guess objective in %r (trees=%r, ntrees=%r, num_class=%r)"
+                    "." % (objs, trees, ntrees, kwargs['num_class']))
         else:
-            kwargs['num_class'] = 1
             kwargs["objective"] = "binary:logistic"
 
     if 'base_score' not in kwargs:

onnxmltools/convert/xgboost/operator_converters/XGBoost.py

@@ -38,7 +38,7 @@ def common_members(xgb_node, inputs):
         booster = xgb_node.get_booster()
         # The json format was available in October 2017.
         # XGBoost 0.7 was the first version released with it.
-        js_tree_list = booster.get_dump(with_stats=True, dump_format = 'json')
+        js_tree_list = booster.get_dump(with_stats=True, dump_format='json')
         js_trees = [json.loads(s) for s in js_tree_list]
         return objective, base_score, js_trees
 

tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py

@@ -20,6 +20,8 @@
 
 
 TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version())
+# PyTorch 1.8.1 supports up to opset version 13.
+HUMMINGBIRD_TARGET_OPSET = min(TARGET_OPSET, 13)
 
 
 class TestLightGbmTreeEnsembleModelsHummingBird(unittest.TestCase):
@@ -46,7 +48,7 @@ def test_lightgbm_booster_classifier(self):
                                data)
         model_onnx, prefix = convert_model(model, 'tree-based classifier',
                                            [('input', FloatTensorType([None, 2]))], without_onnx_ml=True,
-                                           target_opset=TARGET_OPSET,
+                                           target_opset=HUMMINGBIRD_TARGET_OPSET,
                                            zipmap=False)
         dump_data_and_model(X, model, model_onnx,
                             allow_failure="StrictVersion(onnx.__version__) < StrictVersion('1.3.0')",
@@ -63,16 +65,16 @@ def test_lightgbm_booster_classifier_zipmap(self):
                                data)
         model_onnx, prefix = convert_model(model, 'tree-based classifier',
                                            [('input', FloatTensorType([None, 2]))], without_onnx_ml=False,
-                                           target_opset=TARGET_OPSET)
+                                           target_opset=HUMMINGBIRD_TARGET_OPSET)
         assert "zipmap" in str(model_onnx).lower()
         with self.assertRaises(NotImplementedError):
             convert_model(model, 'tree-based classifier',
                                            [('input', FloatTensorType([None, 2]))], without_onnx_ml=True,
-                                           target_opset=TARGET_OPSET)
-                                           
+                                           target_opset=HUMMINGBIRD_TARGET_OPSET)
+
         model_onnx, prefix = convert_model(model, 'tree-based classifier',
                                            [('input', FloatTensorType([None, 2]))], without_onnx_ml=True,
-                                           target_opset=TARGET_OPSET, zipmap=False)
+                                           target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False)
         dump_data_and_model(X, model, model_onnx,
                             allow_failure="StrictVersion(onnx.__version__) < StrictVersion('1.3.0')",
                             basename=prefix + "BoosterBin" + model.__class__.__name__)
@@ -88,7 +90,7 @@ def test_lightgbm_booster_multi_classifier(self):
                                data)
         model_onnx, prefix = convert_model(model, 'tree-based classifier',
                                            [('input', FloatTensorType([None, 2]))], without_onnx_ml=True,
-                                           target_opset=TARGET_OPSET, zipmap=False)
+                                           target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False)
         dump_data_and_model(X, model, model_onnx,
                             allow_failure="StrictVersion(onnx.__version__) < StrictVersion('1.3.0')",
                             basename=prefix + "BoosterBin" + model.__class__.__name__)
@@ -108,7 +110,7 @@ def test_lightgbm_booster_regressor(self):
                                data)
         model_onnx, prefix = convert_model(model, 'tree-based binary regressor',
                                            [('input', FloatTensorType([None, 2]))], without_onnx_ml=True,
-                                           target_opset=TARGET_OPSET, zipmap=False)
+                                           target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False)
         dump_data_and_model(X, model, model_onnx,
                             allow_failure="StrictVersion(onnx.__version__) < StrictVersion('1.0.0')",
                             basename=prefix + "BoosterBin" + model.__class__.__name__)
@@ -203,7 +205,7 @@ def _test_lightgbm_booster_regressor(self):
         y = [0, 1, 1.1]
         data = lightgbm.Dataset(X, label=y)
         model = lightgbm.train(
-            {"boosting_type": "gbdt", "objective": "regression", "n_estimators": 3, 
+            {"boosting_type": "gbdt", "objective": "regression", "n_estimators": 3,
              "min_child_samples": 1, "max_depth": 1, 'num_thread': 1},
             data,
         )

tests/xgboost/test_xgboost_converters.py

@@ -11,7 +11,7 @@
 from sklearn.datasets import (
     load_diabetes, load_iris, make_classification, load_digits)
 from sklearn.model_selection import train_test_split
-from xgboost import XGBRegressor, XGBClassifier, train, DMatrix
+from xgboost import XGBRegressor, XGBClassifier, train, DMatrix, Booster, train as train_xgb
 from sklearn.preprocessing import StandardScaler
 from onnx.defs import onnx_opset_version
 from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER
@@ -181,7 +181,7 @@ def test_xgboost_booster_classifier_multiclass_softmax(self):
                                    random_state=42, n_informative=3)
         x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5,
                                                        random_state=42)
-        
+
         data = DMatrix(x_train, label=y_train)
         model = train({'objective': 'multi:softmax',
                        'n_estimators': 3, 'min_child_samples': 1,
@@ -303,6 +303,44 @@ def test_xgb_empty_tree(self):
         assert_almost_equal(xgb.predict_proba(X), res[1])
         assert_almost_equal(xgb.predict(X), res[0])
 
+    def test_xgb_best_tree_limit(self):
+
+        # Train
+        iris = load_iris()
+        X, y = iris.data, iris.target
+        X_train, X_test, y_train, y_test = train_test_split(X, y)
+        dtrain = DMatrix(X_train, label=y_train)
+        dtest = DMatrix(X_test)
+        param = {'objective': 'multi:softmax', 'num_class': 3}
+        bst_original = train_xgb(param, dtrain, 10)
+        initial_type = [('float_input', FloatTensorType([None, 4]))]
+        bst_original.save_model('model.json')
+
+        onx_loaded = convert_xgboost(
+            bst_original, initial_types=initial_type,
+            target_opset=TARGET_OPSET)
+        sess = InferenceSession(onx_loaded.SerializeToString())
+        res = sess.run(None, {'float_input': X_test.astype(np.float32)})
+        assert_almost_equal(bst_original.predict(dtest, output_margin=True), res[1], decimal=5)
+        assert_almost_equal(bst_original.predict(dtest), res[0])
+
+        # After being restored, the loaded booster is not exactly the same
+        # in memory. `best_ntree_limit` is not saved during `save_model`.
+        bst_loaded = Booster()
+        bst_loaded.load_model('model.json')
+        bst_loaded.save_model('model2.json')
+        assert_almost_equal(bst_loaded.predict(dtest, output_margin=True),
+                            bst_original.predict(dtest, output_margin=True), decimal=5)
+        assert_almost_equal(bst_loaded.predict(dtest), bst_original.predict(dtest))
+
+        onx_loaded = convert_xgboost(
+            bst_loaded, initial_types=initial_type,
+            target_opset=TARGET_OPSET)
+        sess = InferenceSession(onx_loaded.SerializeToString())
+        res = sess.run(None, {'float_input': X_test.astype(np.float32)})
+        assert_almost_equal(bst_loaded.predict(dtest, output_margin=True), res[1], decimal=5)
+        assert_almost_equal(bst_loaded.predict(dtest), res[0])
+
 
 if __name__ == "__main__":
     unittest.main()