Fixes #396, xgboost converter for xgboost >= 1.0.2 (#397)

onnx · Jun 1, 2020 · 18c4dc1 · 18c4dc1
1 parent 4e7c0f0
commit 18c4dc1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/onnxmltools/convert/xgboost/common.py b/onnxmltools/convert/xgboost/common.py
@@ -12,5 +12,9 @@ def get_xgb_params(xgb_node):
     else:
         # XGBoost < 0.7
         params = xgb_node.__dict__
-
+
+    if ('n_estimators' not in params and
+            hasattr(xgb_node, 'n_estimators')):
+        # xgboost >= 1.0.2
+        params['n_estimators'] = xgb_node.n_estimators
     return params        
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -14,4 +14,4 @@ scikit-learn
 scipy
 svm
 wheel
-xgboost<=1.0.2
+xgboost