# -*- coding:UTF-8 -*-
from sklearn.datasets import load_iris
from sklearn import tree
import os
import warnings
warnings.filterwarnings("ignore")
import pydot# need install
print(os.getcwd())
clf=tree.DecisionTreeClassifier(criterion="entropy")
#criterion:gini或者entropy,前者是基尼系数,后者是信息熵。
iris=load_iris()
print(iris.data[0:5])
print(iris.target[0:5])
clf=clf.fit(iris.data,iris.target)
base_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(base_dir, 'tree.dot')
tree.export_graphviz(clf,out_file=file_path)
(graph,)=pydot.graph_from_dot_file(file_path)
graph.write_png(file_path)
结果:
###print(iris.data[0:5])
[[5.1 3.5 1.4 0.2]
[4.9 3. 1.4 0.2]
[4.7 3.2 1.3 0.2]
[4.6 3.1 1.5 0.2]
[5. 3.6 1.4 0.2]]
###print(iris.target[0:5])
[0 0 0 0 0]