如何在 XGBoost 库的 plot_tree 函数中包含特征名称?

FCh*_*Chm 0 python xgboost

我一直在使用XGBoost库来开发二进制分类模型。训练我的模型后,我对可视化单个树感兴趣,以更好地理解我的模型预测。

为此,XGBoost 提供了一个plot_tree函数,但它只显示特征的整数索引。这是我的一棵树的示例:

如何在此图像中包含特征名称而不是特征索引 ( f28)?

FCh*_*Chm 5

plot_treexgboost 中的函数有一个参数fmap,它是“特征映射”文件的路径;这包含特征索引到特征名称的映射。

特征映射文件的文档很少,但它是一个制表符分隔的文件,其中第一列是特征索引(从 0 开始,以特征数量结束),第二列是特征名称,最后一列是显示特征类型的指示器(q=定量特征,i=二元特征)。

一个feature_map.txt文件的例子:

0    feature_name_0    q
1    feature_name_1    i
2    feature_name_2    q
…          …           … 
Run Code Online (Sandbox Code Playgroud)

使用这个制表符分隔的文件,您可以从训练有素的模型实例中绘制树:

import xgboost
model = xgboost.XGBClassifier()

# train the model
model.fit(X, y)

# plot the decision tree, providing path to feature map file

xgboost.plot_tree(model,  num_trees=0, fmap='feature_map.txt')
Run Code Online (Sandbox Code Playgroud)

使用此函数显示绘图: