决策树的使用案例

import numpy as np

from sklearn.tree import DecisionTreeClassifier

from sklearn import datasets

from sklearn.model_selection import train_test_split

from sklearn import tree

import matplotlib.pyplot as plt

1. 构建数据集

iris = datasets.load_iris()

X = iris['data']
y = iris['target']
# 样本的属性名
fn = iris['feature_names']

# 分为训练数据和测试数据
X_train,X_test,y_train,y_test = train_test_split(X,y)

2. 决策树分类

# 创建模型，max_depth参数为决策树的深度
model = DecisionTreeClassifier(max_depth=None,criterion='entropy')
# 训练模型
model.fit(X_train,y_train)
# 模型预测的结果
y_ = model.predict(X_test)
print('真实类别：',y_test)
print('算法预测类别：',y_)
print('准确率：',model.score(X_test,y_test))

结果为

3. 绘制决策树

# 设置图片尺寸
plt.figure(figsize=(24,24))
# 绘制决策树
_ = tree.plot_tree(model,filled = True,feature_names=fn)

结果为

4. 决策树剪枝

通过设置max_depth来调整树深度，即剪枝操作

# 只修改此行代码
model = DecisionTreeClassifier(max_depth=2,criterion='entropy')

结果为

可以看到此时决策树完成了剪枝操作。