Support objective='reg:logistic' for XGBoost converters (#270)

Support objective='reg:logistic' for XGBoost converters

Author: xadupre

onnxmltools/convert/xgboost/operator_converters/XGBoost.py

@@ -219,10 +219,10 @@ def convert(scope, operator, container):
         objective, base_score, js_trees = XGBConverter.common_members(xgb_node, inputs)
 
         params = XGBConverter.get_xgb_params(xgb_node)
-        
+
         attr_pairs = XGBClassifierConverter._get_default_tree_attribute_pairs()
         XGBConverter.fill_tree_attributes(js_trees, attr_pairs, [1 for _ in js_trees], True)
-        
+
         if len(attr_pairs['class_treeids']) == 0:
             raise RuntimeError("XGBoost model is empty.")
         ncl = (max(attr_pairs['class_treeids']) + 1) // params['n_estimators']
@@ -251,6 +251,13 @@ def convert(scope, operator, container):
             container.add_node('TreeEnsembleClassifier', operator.input_full_names,
                                operator.output_full_names,
                                op_domain='ai.onnx.ml', **attr_pairs)
+        elif objective == "reg:logistic":
+            ncl = len(js_trees) // params['n_estimators']
+            if ncl == 1:
+                ncl = 2
+            container.add_node('TreeEnsembleClassifier', operator.input_full_names,
+                               operator.output_full_names,
+                               op_domain='ai.onnx.ml', **attr_pairs)
         else:
             raise RuntimeError("Unexpected objective: {0}".format(objective))            
 

onnxmltools/convert/xgboost/shape_calculators/Classifier.py

@@ -26,6 +26,8 @@ def calculate_xgboost_classifier_output_shapes(operator):
         ncl = 2
     else:
         ncl = ntrees // params['n_estimators']
+        if objective == "reg:logistic" and ncl == 1:
+            ncl = 2
     operator.outputs[0].type = Int64TensorType(shape=[N])
     operator.outputs[1].type = operator.outputs[1].type = FloatTensorType([N, ncl])
 

tests/xgboost/test_xgboost_converters.py

@@ -12,7 +12,7 @@
 
 class TestXGBoostModels(unittest.TestCase):
 
-    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converted not tested on python 2")
+    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converter not tested on python 2")
     def test_xgb_regressor(self):
         iris = load_iris()
         X = iris.data[:, :2]
@@ -24,7 +24,7 @@ def test_xgb_regressor(self):
         self.assertTrue(conv_model is not None)
         dump_single_regression(xgb, suffix="-Dec4")
 
-    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converted not tested on python 2")
+    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converter not tested on python 2")
     def test_xgb_classifier(self):
         iris = load_iris()
         X = iris.data[:, :2]
@@ -35,9 +35,9 @@ def test_xgb_classifier(self):
         xgb.fit(X, y)
         conv_model = convert_xgboost(xgb, initial_types=[('input', FloatTensorType(shape=[1, 'None']))])
         self.assertTrue(conv_model is not None)
-        dump_binary_classification(xgb, verbose=True)
+        dump_binary_classification(xgb)
 
-    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converted not tested on python 2")
+    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converter not tested on python 2")
     def test_xgb_classifier_multi(self):
         iris = load_iris()
         X = iris.data[:, :2]
@@ -49,6 +49,32 @@ def test_xgb_classifier_multi(self):
         self.assertTrue(conv_model is not None)
         dump_multiple_classification(xgb, allow_failure="StrictVersion(onnx.__version__) < StrictVersion('1.3.0')")
 
+    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converter not tested on python 2")
+    def test_xgb_classifier_multi_reglog(self):
+        iris = load_iris()
+        X = iris.data[:, :2]
+        y = iris.target
+
+        xgb = XGBClassifier(objective='reg:logistic')
+        xgb.fit(X, y)
+        conv_model = convert_xgboost(xgb, initial_types=[('input', FloatTensorType(shape=[1, 2]))])
+        self.assertTrue(conv_model is not None)
+        dump_multiple_classification(xgb, suffix="RegLog",
+                                     allow_failure="StrictVersion(onnx.__version__) < StrictVersion('1.3.0')")
+
+    @unittest.skipIf(sys.version_info[0] == 2, reason="xgboost converter not tested on python 2")
+    def test_xgb_classifier_reglog(self):
+        iris = load_iris()
+        X = iris.data[:, :2]
+        y = iris.target
+        y[y == 2] = 0
+
+        xgb = XGBClassifier(objective='reg:logistic')
+        xgb.fit(X, y)
+        conv_model = convert_xgboost(xgb, initial_types=[('input', FloatTensorType(shape=[1, 2]))])
+        self.assertTrue(conv_model is not None)
+        dump_binary_classification(xgb, suffix="RegLog")
+
 
 if __name__ == "__main__":
     unittest.main()