XGBoost 模型 - 绘制决策树

xgboost画图时遇到如下若干坑

1. 图像过小，看不清内容

2. 只显示特征编号，不显示特征名

3. 怎么把图像保存

 

目录

解决方法：

使用xgboost自带的plot_tree函数绘制决策树

绘制出来的图形不显示特征名，只是显示特征序号

调整Figure大小

保存图像文件

---

 

 

解决方法：

使用xgboost自带的plot_tree函数绘制决策树

训练模型xgb_clf

from xgboost.sklearn import XGBClassifier
xgb_clf = XGBClassifier()
xgb_clf .fit(x_train, y_train)

调用plot_tree函数

from xgboost import plot_tree
plot_tree(xgb_clf)

 

然而绘制出来的图形不显示特征名，只是显示特征序号

查看原函数plot_tree(booster, fmap='', num_trees=0, rankdir='UT', ax=None, **kwargs)发现一参数fmap不知何意。 
几经尝试，最后发现kaggle大神已给出解决方案：https://www.kaggle.com/mmueller/xgb-feature-importance-python

def ceate_feature_map(features):
    outfile = open('xgb.fmap', 'w')
    i = 0
    for feat in features:
        outfile.write('{0}\t{1}\tq\n'.format(i, feat))
        i = i + 1
    outfile.close()
ceate_feature_map(train_data.columns)#特征名列表

 

调整Figure大小

plot_tree未提供修改图像大小的参数，这里直接通过在新建的Figure,Axes对象，调整Figure大小，再在其上画决策树图的方法实现调整大小

fig, ax = plt.subplots()
fig.set_size_inches(60, 30)
plot_tree(xgb_clf, ax=ax, fmap='xgb.fmap')

保存图像文件

通过fig的savefig方法来保存图像

fig.savefig('xgb_tree.png')