如何在scikit学习中将特征名称添加到决策树的输出中 [英] How to add feature names to output of decision tree in scikit learn

问题描述

我正在 scikit-learn 中构建决策树，然后想生成树的 pdf.我输出树的工作流程大致如下.

I am building a decision tree in scikit-learn then want to produce a pdf of the tree. My workflow to output the tree is roughly as follows.

vec = DictVectorizer()
data_vectorized = vec.fit_transform(data)
vec.get_feature_names() #Shows feature names

clf = tree.DecisionTreeClassifier()
clf = clf.fit(data_vectorized, Labels)

from sklearn.externals.six import StringIO  
import pydot 
dot_data = StringIO() 
tree.export_graphviz(clf, out_file=dot_data) 
graph = pydot.graph_from_dot_data(dot_data.getvalue()) 
graph.write_pdf("tree.pdf") 

然而，这个 pdf 的每个节点都显示了使用 data_vectorized[i] 对某些 i 的比较.当 data_vectorized 大而稀疏时，这很难解释.

However each node of this pdf shows a comparison using data_vectorized[i] for some i. This is quite hard to interpret when data_vectorized is large and sparse.

我怎样才能让它显示功能的名称?

How can I get it show the name of the feature instead?

此图像显示了您得到的示例(当特征位于变量 X 中时.例如，我希望 X[2] 被特征名称替换.

This image shows an example of what you get (when the features are in a variable X. I would like X[2], for example, to be replaced by the name of the feature.

推荐答案

尝试将导出更改为:

tree.export_graphviz(clf, out_file=dot_data, feature_names=vec.get_feature_names()) 

这篇关于如何在scikit学习中将特征名称添加到决策树的输出中的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！