本文基于经典的 Iris(鸢尾花)数据集,使用 scikit-learn 中的决策树完成一个完整的三分类任务,覆盖数据加载 → 数据探查 → 特征工程 → 模型训练 → 模型评估 → 可视化的完整机器学习流程,非常适合机器学习初学者上手。文末附完整可运行代码以及几个初学者容易踩的坑。
一、简述
鸢尾花数据集(Iris Dataset)是机器学习领域的 “Hello World”,由统计学家 Fisher 于 1936 年整理。数据集一共记录了 150 朵鸢尾花的 4 个形态学特征,目标是根据这些特征判断它属于以下 3 个品种中的哪一种:
山鸢尾(Setosa)
变色鸢尾(Versicolor)
维吉尼亚鸢尾(Virginica)
4 个特征分别是:
特征 含义
sepal length (cm) 花萼长度
sepal width (cm) 花萼宽度
petal length (cm) 花瓣长度
petal width (cm) 花瓣宽度

本文选用决策树(Decision Tree) 作为分类器。它原理直观、可解释性强、对数据量要求低,而且训练完还能输出"哪个特征最重要",非常适合入门理解机器学习到底在学什么。

二、环境准备
需要用到以下第三方库,如果没装可以先安装:

pip install pandas scikit-learn matplotlib seaborn

andas:数据处理
scikit-learn:数据集、模型与评估指标
matplotlib + seaborn:可视化

三、完整流程实现

  1. 导入依赖
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

# 设置中文字体,防止图表中文乱码
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei"]
plt.rcParams["axes.unicode_minus"] = False

⚠️ 坑点 1:在 Windows 上画图如果出现中文方块乱码,务必加上上面两行字体配置,注意是 Microsoft YaHei,不要写成 Hicrosoft。

  1. 加载数据集
    scikit-learn 内置了鸢尾花数据集,一行代码即可加载:
ris = load_iris()
X, y = iris.data, iris.target                 # 特征矩阵 X,标签 y
feature_names = iris.feature_names            # 特征名
target_names = ["山鸢尾", "变色鸢尾", "维吉尼亚鸢尾"]  # 用中文名替换默认英文标签

# 转成 DataFrame 方便查看
df = pd.DataFrame(X, columns=feature_names)
df["target"] = y
  1. 数据探查(EDA)
    建模前先了解数据,是一个好习惯:
print("数据基本信息:")
df.info()

print("\n数据前 5 行:")
print(df.head())

print("\n各类别样本数量:")
print(df["target"].value_counts().sort_index())

运行后可以看到,数据集共 150 行、5 列(4 个特征 + 1 个标签) ,且没有缺失值,三个类别各 50 条,是一个完全平衡的数据集:

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 150 entries, 0 to 149
Data columns (total 5 columns):
 #   Column             Non-Null Count  Dtype
---  ------             --------------  -----
 0   sepal length (cm)  150 non-null    float64
 1   sepal width (cm)   150 non-null    float64
 2   petal length (cm)  150 non-null    float64
 3   petal width (cm)   150 non-null    float64
 4   target             150 non-null    int64
dtypes: float64(4), int64(1)

重复值检查

duplicate_rows = df[df.duplicated()]
print(f"重复样本数:{len(duplicate_rows)}")

重复样本数:1

数据中存在 1 条重复记录。由于数量极少且本文重点是走通流程,这里仅做提示、暂不删除;实际工程中可结合业务决定是否去重。
异常值检查(IQR 法)
以"花瓣长度"为例,用四分位距(IQR)方法检测异常值:

1 = df["petal length (cm)"].quantile(0.25)
q3 = df["petal length (cm)"].quantile(0.75)
iqr = q3 - q1
lower_bound = q1 - 1.5 * iqr
upper_bound = q3 + 1.5 * iqr

outliers = df[
    (df["petal length (cm)"] < lower_bound) |
    (df["petal length (cm)"] > upper_bound)
]
print(f"异常值样本数:{len(outliers)}")

异常值样本数:0

花瓣长度没有 IQR 异常值,说明该特征数据质量良好。
4. 特征标准化 & 划分训练集 / 测试集

# 标准化:均值为 0,方差为 1
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 按 8:2 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=42
)

print(f"训练集大小:{X_train.shape},测试集大小:{X_test.shape}")

训练集大小:(120, 4),测试集大小:(30, 4)

💡 值得思考的一点:标准化对决策树其实几乎没有影响。因为决策树是通过比较"特征值与阈值的大小关系"来分裂节点的,单调的缩放和平移不会改变这种大小关系。这里保留标准化步骤,主要是为了演示一套通用的机器学习流程(换成 KNN、SVM、逻辑回归等对尺度敏感的模型时就非常必要了)。
5. 训练决策树模型

model = DecisionTreeClassifier(max_depth=3, random_state=42)
model.fit(X_train, y_train)

这里设置 max_depth=3 限制树的最大深度,既能防止模型过拟合,也方便后续可视化和理解;random_state 固定随机种子,保证结果可复现。
6. 查看特征重要性
决策树的一大优点是可解释,训练后可以直接查看每个特征对分类的贡献:

feature_importance = pd.DataFrame({
    "特征": feature_names,
    "重要性": model.feature_importances_
}).sort_values("重要性", ascending=False)

print("特征重要性排序:")
print(feature_importance)

特征重要性排序:
                 特征      重要性
2  petal length (cm)  0.934626
3   petal width (cm)  0.065374
1   sepal width (cm)  0.000000
0  sepal length (cm)  0.000000

可以非常直观地看出:花瓣长度(petal length)贡献了约 93.5% 的区分能力,花瓣宽度次之,而花萼相关特征在深度为 3 时几乎没被用到。这也符合植物学常识——三个品种的花瓣差异远比花萼差异明显。
还可以画成柱状图:

plt.figure(figsize=(8, 5))
sns.barplot(data=feature_importance, x="重要性", y="特征", palette="viridis")
plt.title("决策树特征重要性")
plt.tight_layout()
plt.show()

  1. 模型预测与评估
y_pred = model.predict(X_test)

# 准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率:{accuracy:.4f}")

# 精确率 / 召回率 / F1
print("\n分类报告:")
print(classification_report(y_test, y_pred, target_names=target_names, digits=4))
模型准确率:1.0000

分类报告:
            precision  recall  f1-score  support
      山鸢尾     1.0000  1.0000   1.0000       10
     变色鸢尾    1.0000  1.0000   1.0000        9
   维吉尼亚鸢尾  1.0000  1.0000   1.0000       11

    accuracy                         1.0000       30
   macro avg    1.0000  1.0000   1.0000       30
weighted avg   1.0000  1.0000   1.0000       30

在 30 条测试样本上,模型取得了 100% 的准确率。
8. 混淆矩阵可视化
混淆矩阵能清楚展示每个类别被预测正确 / 错误的情况:

cm = confusion_matrix(y_test, y_pred)

plt.figure(figsize=(6, 5))
sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
            xticklabels=target_names,
            yticklabels=target_names)
plt.xlabel("预测类别")
plt.ylabel("真实类别")
plt.title("混淆矩阵")
plt.tight_layout()
plt.show()

[[10  0  0]
 [ 0  9  0]
 [ 0  0 11]]

混淆矩阵是一条完美的对角线,说明三类样本全部预测正确,没有任何误分。
⚠️ 坑点 2(最容易报错的一行) :混淆矩阵一定要用 sns.heatmap(cm, annot=True, fmt=“d”, …) 来画。
annot=True 表示在格子里显示数值;
fmt=“d” 表示按整数格式显示(注意是 fmt 不是 fnt,漏写会导致数字以科学计数法或报错显示);
如果误写成 plt.figure(cm, annot=True, fnt=‘d’, …) 会直接报 TypeError,因为这些是 seaborn 热力图的参数,不是 matplotlib figure 的参数。
四、(可选)可视化决策树结构
想进一步理解决策树"怎么判断的",可以把整棵树画出来:

from sklearn.tree import plot_tree

plt.figure(figsize=(14, 8))
plot_tree(
    model,
    feature_names=feature_names,
    class_names=target_names,
    filled=True,
    rounded=True,
    fontsize=10
)
plt.title("决策树结构(max_depth=3)")
plt.tight_layout()
plt.show()

图中每个节点会显示判断条件(如 petal length <= 阈值)、不纯度(gini)、样本数和分类结果,可以清楚看到模型是如何一步步把三种花区分开的。
五、完整代码

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

# ---------- 全局配置 ----------
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei"]
plt.rcParams["axes.unicode_minus"] = False

# ---------- 1. 加载数据 ----------
iris = load_iris()
X, y = iris.data, iris.target
feature_names = iris.feature_names
target_names = ["山鸢尾", "变色鸢尾", "维吉尼亚鸢尾"]

df = pd.DataFrame(X, columns=feature_names)
df["target"] = y

# ---------- 2. 数据探查 ----------
print("数据基本信息:")
df.info()
print("\n数据前 5 行:")
print(df.head())
print("\n各类别样本数:")
print(df["target"].value_counts().sort_index())

# 重复值
duplicate_rows = df[df.duplicated()]
print(f"\n重复样本数:{len(duplicate_rows)}")

# 异常值(IQR)
q1 = df["petal length (cm)"].quantile(0.25)
q3 = df["petal length (cm)"].quantile(0.75)
iqr = q3 - q1
lower_bound = q1 - 1.5 * iqr
upper_bound = q3 + 1.5 * iqr
outliers = df[(df["petal length (cm)"] < lower_bound) |
              (df["petal length (cm)"] > upper_bound)]
print(f"异常值样本数:{len(outliers)}")

# ---------- 3. 标准化 + 划分数据集 ----------
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=42
)
print(f"\n训练集大小:{X_train.shape},测试集大小:{X_test.shape}")

# ---------- 4. 训练模型 ----------
model = DecisionTreeClassifier(max_depth=3, random_state=42)
model.fit(X_train, y_train)

# ---------- 5. 特征重要性 ----------
feature_importance = pd.DataFrame({
    "特征": feature_names,
    "重要性": model.feature_importances_
}).sort_values("重要性", ascending=False)
print("\n特征重要性排序:")
print(feature_importance)

# ---------- 6. 预测与评估 ----------
y_pred = model.predict(X_test)
print(f"\n模型准确率:{accuracy_score(y_test, y_pred):.4f}")
print("\n分类报告:")
print(classification_report(y_test, y_pred,
                            target_names=target_names, digits=4))

# ---------- 7. 混淆矩阵 ----------
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(6, 5))
sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
            xticklabels=target_names, yticklabels=target_names)
plt.xlabel("预测类别")
plt.ylabel("真实类别")
plt.title("混淆矩阵")
plt.tight_layout()
plt.show()

# ---------- 8.(可选)决策树结构 ----------
plt.figure(figsize=(14, 8))
plot_tree(model, feature_names=feature_names, class_names=target_names,
          filled=True, rounded=True, fontsize=10)
plt.title("决策树结构(max_depth=3)")
plt.tight_layout()
plt.show()

六、结果分析与小结
数据层面:Iris 数据集干净、平衡(150 条、三类各 50 条、无缺失值),只有 1 条重复记录,花瓣长度无 IQR 异常值。
特征层面:决策树认为花瓣长度 / 宽度是区分品种的关键特征(花瓣长度重要性约 93.5%),花萼特征几乎不起作用。
模型层面:限制深度为 3 的决策树在测试集上达到 100% 准确率,三类样本全部正确分类。
需要提醒的是,100% 的准确率并不代表模型"天下无敌",而是因为 Iris 数据集本身类别可分性极高、且数据量很小。在更复杂的真实数据上,通常还需要配合交叉验证、网格搜索调参、以及学习曲线等手段来评估模型的泛化能力。
初学者可继续尝试的方向
用 export_text(model) 或 plot_tree 输出决策规则,加深对树模型的理解;
调整 max_depth、criterion(gini / entropy)等超参数,观察准确率变化;
改用 KNN、SVM、逻辑回归、随机森林等模型做横向对比;
使用 cross_val_score 做 5 折交叉验证,让评估结果更稳健。

更多推荐