sklearn是一个机器学习的python库,其中可以很方便的进行决策树算法的进行,但是没有内置的api可以进行决策树的可视化操作,可以考虑用另一个第三方库graphviz进行可视化的操作,使用pip install graphviz安装库,在graphviz官网上下载安装包进行安装(http://www.graphviz.org/)并配置环境变量。
决策树算法的实现:
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier,export_graphviz
import graphviz
iris = load_iris()
x_train,x_test,y_train,y_test=train_test_split(iris.data,iris.target,random_state=22)
#决策树预估器
estimator=DecisionTreeClassifier(criterion="entropy")
estimator.fit(x_train,y_train)
#模型评估
y_predict=estimator.predict(x_test)
print("预测结果:\n",y_predict)
print("比对结果:\n",y_predict==y_test)
score=estimator.score(x_test,y_test)
print("准确率为:\n",score)
dot_data=export_graphviz(estimator,out_file=None,feature_names=iris.feature_names,class_names=iris.target_names)
其中可以通过以下语句生成.dot文件:
export_graphviz(estimator,out_file="iris_tree.dot",feature_names=iris.feature_names)
可以通过以下代码生成图片:
graph=graphviz.Source(dot_data)
graph
若要是想生成图片可以使用cmd命令(前提是graphviz已经配置在环境变量中)
例如我已经生成了iris_tree.dot文件是使用以下命令:
可以生成tree.png文件
也可以通过:
graph.render('test.gv',view=True)
生成pdf文件。