take best tree if possible

Author: erdogant

treeplot/examples.py

@@ -30,6 +30,13 @@
 
 ax = treeplot.plot(model_gradientboost)
 
+# %% XGBOOST EXAMPLE
+import xgboost as xgb
+model_xgb = xgb.XGBClassifier(n_estimators=100, max_depth=2, random_state=0).fit(X, y)
+
+ax = treeplot.plot(model_xgb)
+ax = treeplot.xgboost(model_xgb, plottype='vertical')
+
 # %% XGBOOST EXAMPLE
 from xgboost import XGBClassifier
 model_xgb = XGBClassifier(n_estimators=100, max_depth=2, random_state=0).fit(X, y)

treeplot/treeplot.py

@@ -101,7 +101,10 @@ def xgboost(model, featnames=None, num_trees=None, plottype='horizontal', figsiz
 
     if plottype=='horizontal': plottype='UD'
     if plottype=='vertical': plottype='LR'
-    if num_trees is None: num_trees = model.best_iteration
+    if (num_trees is None) and hasattr(model, 'best_iteration'):
+        num_trees = model.best_iteration
+    elif num_trees is None:
+        num_trees = 0
 
     try:
         fig, ax = plt.subplots(1, 1, figsize=figsize)