《scikit-learn》决策树之鸢尾花分类

有了上一博客的基础，我们来看看怎么操作鸢尾花的分裂问题。也是做一个简单的总结和回顾。

直接整代码了。

from sklearn import tree
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import graphviz

# 分类问题举例，加载自带的鸢尾花数据集，150个样本，4个属性，3个类别。
iris = load_iris()
print(iris.data.shape)
print(iris.target.shape)
print(iris.target_names)  # 所有类别标签地名称
print(iris.feature_names)  # 所有特征属性的名称

# 把红酒数据进行切分，切分成训练集和测试集合，切分比例一般是7:3
data_train, data_test, target_train, target_test = train_test_split(iris.data, iris.target, test_size=0.3)
print(data_train.shape)
print(target_train.shape)
print(data_test.shape)
print(target_test.shape)


# 进行模型的搭建和训练
clf = tree.DecisionTreeClassifier(criterion='entropy')  # 定义一个决策树的实例，决策树节点的分裂选择的是 根据信息熵划分
clf = clf.fit(data_train, target_train)  # 用训练数据和标签，进行训练。简单的就这么简单就训练出了模型。
score = clf.score(data_test, target_test)  # 用已经训练好的模型，对测试数据进行预测准确率，0~1，越接近1越好。
print(score)


# 把这棵树clf进行图像化，让人看着更加清晰
dot_data = tree.export_graphviz(clf,  # 决策树本身
                                feature_names=iris.feature_names,  # 特征名称
                                class_names=iris.target_names,  # 类别名称
                                filled=True,  # 给图形填充颜色
                                rounded=True  # 图形的节点是圆角矩形
                                )
graph = graphviz.Source(dot_data)
graph.view(filename='kk.png', cleanup=True)

效果如下：