AI人工智能 决策树分类器
决策树是一种二叉树流程图,每个节点根据某个特征变量对样本集进行划分,最终通过叶子节点输出分类结果。本文将构建基于“身高”和“头发长度”的性别预测分类器。
1. 前置条件
需安装pydotplus和graphviz库(用于决策树可视化):
graphviz:通过dot文件绘制图形的工具;pydotplus:操作Graphviz的Dot语言的Python模块。 安装命令:pip install pydotplus graphviz。
2. 实现步骤
步骤1:导入库
import pydotplus
from sklearn import tree
from sklearn.metrics import classification_report
from sklearn import cross_validation
import collections
步骤2:定义数据集
特征为“身高(cm)”和“头发长度(cm)”,标签为“Man(男性)”或“Woman(女性)”:
## 特征:[身高, 头发长度]
X = [[165, 19], [175, 32], [136, 35], [174, 65], [141, 28], [176, 15],
[131, 32], [166, 6], [128, 32], [179, 10], [136, 34], [186, 2],
[126, 25], [176, 28], [112, 38], [169, 9], [171, 36], [116, 25], [196, 25]]
## 标签:性别
Y = ['Man', 'Woman', 'Woman', 'Man', 'Woman', 'Man', 'Woman', 'Man', 'Woman',
'Man', 'Woman', 'Man', 'Woman', 'Woman', 'Woman', 'Man', 'Woman', 'Woman', 'Man']
## 特征名
data_feature_names = ['height(身高)', 'length of hair(头发长度)']
步骤3:划分训练集与测试集
X_train, X_test, Y_train, Y_test = cross_validation.train_test_split(
X, Y, test_size=0.40, random_state=5
)
步骤4:构建并训练决策树
## 初始化决策树分类器
clf = tree.DecisionTreeClassifier()
## 训练模型(使用全量数据,也可替换为X_train和Y_train)
clf = clf.fit(X, Y)
步骤5:模型预测
预测身高133cm、头发长度37cm的样本性别:
prediction = clf.predict([[133, 37]])
print(prediction) # 输出:['Woman'](女性)
步骤6:决策树可视化
将决策树保存为PNG图片:
## 导出决策树的dot格式数据
dot_data = tree.export_graphviz(
clf, feature_names=data_feature_names, out_file=None,
filled=True, rounded=True
)
## 生成图形对象
graph = pydotplus.graph_from_dot_data(dot_data)
## 设置节点颜色
colors = ('orange', 'yellow')
edges = collections.defaultdict(list)
for edge in graph.get_edge_list():
edges[edge.get_source()].append(int(edge.get_destination()))
for edge in edges:
edges[edge].sort()
for i in range(2):
dest = graph.get_node(str(edges[edge][i]))[0]
dest.set_fillcolor(colors[i])
## 保存为PNG图片
graph.write_png('Decisiontree16.png')
决策树核心划分规则
- 根节点:头发长度≤22.0cm(基尼系数=0.488,19个样本:8男11女);
- 左子树(头发长度≤22.0cm):6个样本全为男性(基尼系数=0.0);
- 右子树(头发长度>22.0cm):13个样本(2男11女),进入下一级划分;
- 次级节点:身高≤186.0cm(基尼系数=0.153,12个样本:1男11女);
- 左子树(身高≤186.0cm):11个样本全为女性(基尼系数=0.0);
- 右子树(身高>186.0cm):1个样本为男性(基尼系数=0.0)。

可以改变预测中的特征值来测试它。