如何画XGBoost里面的决策树

update:
在 xgboost的新版本(0.8)中,已无需生成 fmap来获取列名。

xgboost画图时遇到如下若干坑

1. 图像过小，看不清内容
2. 只显示特征编号，不显示特征名
3. 怎么把图像保存

1 更新-直接显示矢量图

xgboost的新版本(0.8)对pandas的DataFrame支持已经很完善了，我们可以使用xgboost的sklearn 接口，直接训练DataFrame数据，同时也可以将列名直接作为特征名而无需再生成fmap文件。(也可以手动修改xgb.Booster.feature_names = ['col1','col2',])
【注意！】因为xgboost内部会根据columns自动生成fmap，而fmap文件是以空格分隔每一项的，因此， 列名中不能出现空格！
在jupyter notebook中也可以使用xgb.to_graph方法来通过graphviz绘制后直接显示出来，该方法可显示矢量图，并可配置图字体的大小。

这里我们使用breast_cancer数据集来测试：

1.1 替换列名中的空格

import xgboost as xgb
import pandas as pd

from sklearn.datasets import load_breast_cancer
data =load_breast_cancer()

查看发现列名中含有空格

In [2]: data['feature_names'][0]
Out[2]: 'mean radius'

替换空格

# columns name can not have blank if you want to draw tree with feature name
data['feature_names'] = [ ''.join([character if character!=' ' else '_' for character in featureName ]) for featureName in data['feature_names']]

In [4]: data['feature_names'][0]
Out[4]: 'mean_radius'

1.2 使用XGBClassifier训练并画树

df = pd.DataFrame(data['data'],columns=data['feature_names'])
clf = xgb.XGBClassifier()
clf.fit(df,data['target'])
xgb.to_graphviz(clf)

1.3 优化节点形状

xgboost提供的绘图函数，节点shape为circle，画出来的图奇丑无比，我们对节点参数进行修改。
在xgboost源代码的ploting.py文件中修改以下几行：

# line 166
def to_graphviz(booster, fmap='', num_trees=0, rankdir='UT',
                yes_color='#0000FF', no_color='#FF0000',conditionNodeParams={},leafNodeParams={}, **kwargs):
 # line 215
node = _parse_node(graph, text,conditionNodeParams=conditionNodeParams,leafNodeParams= leafNodeParams)
# line 127
def _parse_node(graph, text,conditionNodeParams,leafNodeParams):
# line 132
        graph.node(node, label=match.group(2), **conditionNodeParams)
# line 137
        graph.node(node, label=match.group(2),**leafNodeParams)

修改条件节点的参数如下

cNodeParams = {'shape':'box',
               'style':'filled,rounded',
               'fillcolor':'#78bceb'
              }
lNodeParams = {'shape':'box',
               'style':'filled',
               'fillcolor':'#e48038'
              }
graphParams = {
    'graph_attr':{}
    
}
x = plotting.to_graphviz(clf,conditionNodeParams=cNodeParams,leafNodeParams=lNodeParams,**{'size':str(15)})

得到下面的图。

---

---

2 原方法：

2.1 plot_tree画图

在使用xgboost训练出模型xgbClf后:

import xgboost as xgb
from xgboost.sklearn import XGBClassifier
xgbClf = XGBClassifier()
xgbClf.fit(xTrain,yTrain)

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

xgb.plot_tree(xgbClf)

2.2 生成fmap文件

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

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

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)#特征名列表

该函数生成一fmap文件名为xgb.fmap，那么我们可调用该fmap文件添加特征名。

xgb.plot_tree(xgbClf,fmap='xgb.fmap')

2.3 直接对Figure对象调整大小

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

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

后续若想再次显示图像，直接在jupyter notebook的新建cell里输入：

fig

2.4 保存图像文件

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

fig.savefig('xgb_tree.jpg')

The End