如何在 XGBoost 库的 plot_tree function 中包含特征名称？

Question

我一直在使用XGBoost库来开发二进制分类 model。在训练了我的 model 之后，我有兴趣可视化各个树以更好地理解我的模型预测。

为此，XGBoost 提供了一个plot_tree function，但它只显示了该特征的 integer 索引。 这是我的一棵树的示例：

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

Answer 1

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

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

feature_map.txt文件的示例：

0    feature_name_0    q
1    feature_name_1    i
2    feature_name_2    q
…          …           … 

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

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')

使用此函数显示绘图：

Answer 2

我发现最简单的方法是直接在 model 上设置功能名称。 显然，顺序需要正确。

model = XGBClassifier()    
model.get_booster().feature_names = feature_names # As a list or tuple!

之后，特征名称将与 plot_tree function 一起出现。