from sklearn.tree import DecisionTreeClassifier
from sklearn import tree
from sklearn.externals.six import StringIO
import pydot
clf = DecisionTreeClassifier(max_depth=3)
# features.remove("rule_score")
# features.remove("rules_cnt")
features.remove('type_cnt')
x = ip_features[features]
X = x.values
y = ip_features['level_9'].apply(lambda x: int(x))
Y = y.values
clf.fit(X, Y)
dot_data = StringIO()
tree.export_graphviz(clf, out_file=dot_data, feature_names=features, class_names=['0', '1'])
graph = pydot.graph_from_dot_data(dot_data.getvalue())
graph[0].write_pdf("tree5.pdf")
sklearn训练决策树并可视化
最新推荐文章于 2022-10-02 20:20:08 发布