模型剪枝实战：如何删除冗余结构，让机器学习模型更轻量？

本文通过两个实际案例介绍剪枝技术：使用决策树剪枝改善乳腺癌数据分类，以及使用 PyTorch 对手写数字识别网络进行权重剪枝。 两个案例分别展示剪枝在提高泛化能力和降低有效参数数量方面的应用。

一、什么是剪枝？

模型训练过程中，部分结构可能对最终预测结果贡献较小。

例如，一棵决策树可能拥有上百个判断节点，其中部分分支只用于区分少量特殊样本。

神经网络也可能存在大量权重很小、对预测影响有限的连接。

剪枝的核心思想是：

找到模型中相对不重要的部分，将其删除或置零，再验证模型是否仍能保持较好的预测性能。

常见剪枝方法包括：

需要注意，剪枝的作用取决于模型类型。

决策树剪枝通常侧重减少过拟合和简化规则；神经网络剪枝则更多用于模型压缩和部署优化。

二、案例一：决策树后剪枝，改善分类模型

1. 业务背景

使用 scikit-learn 内置的 Wisconsin Breast Cancer 数据集。

该数据集包含 569 条样本、30 个数值特征，目标是根据细胞核测量结果区分两类诊断标签。

我们先训练一棵完整决策树，再利用成本复杂度剪枝（Cost-Complexity Pruning）删除部分分支。

剪枝的目标函数为：

[
R_\alpha(T)=R(T)+\alpha|T|
]

其中：

• (R(T))：决策树的不纯度代价。
• (|T|)：叶节点数量。
• (\alpha)：剪枝强度。

ccp_alpha 越大，通常意味着剪枝程度越高。

2. Python 完整实践

安装依赖：

pip install scikit-learn numpy

代码如下：

import numpy as np

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score

# 1. 加载数据
X, y = load_breast_cancer(return_X_y=True)

# 2. 划分训练开发集、测试集
X_dev, X_test, y_dev, y_test = train_test_split(
    X, y, test_size=0.2,
    stratify=y, random_state=42
)

# 再划分训练集和验证集
X_train, X_valid, y_train, y_valid = train_test_split(
    X_dev, y_dev,
    test_size=0.25,
    stratify=y_dev,
    random_state=42
)

# 3. 未剪枝模型
base = DecisionTreeClassifier(random_state=42)
base.fit(X_dev, y_dev)

# 4. 生成剪枝候选参数
tree = DecisionTreeClassifier(random_state=42)

path = tree.cost_complexity_pruning_path(
    X_train, y_train
)

# 5. 在验证集上选择最佳剪枝强度
best_alpha = 0
best_acc = -1

for alpha in path.ccp_alphas[:-1]:

    model = DecisionTreeClassifier(
        ccp_alpha=float(alpha),
        random_state=42
    )

    model.fit(X_train, y_train)

    acc = accuracy_score(
        y_valid, model.predict(X_valid)
    )

    if acc > best_acc or (
        acc == best_acc and alpha > best_alpha
    ):
        best_acc = acc
        best_alpha = alpha

# 6. 重新训练剪枝模型
pruned = DecisionTreeClassifier(
    ccp_alpha=float(best_alpha),
    random_state=42
)

pruned.fit(X_dev, y_dev)

# 7. 对比模型
for name, model in [
    ("原始决策树", base),
    ("剪枝决策树", pruned)
]:
    pred = model.predict(X_test)

    print(name)
    print("深度:", model.get_depth())
    print("叶节点:", model.get_n_leaves())
    print("准确率:", accuracy_score(y_test, pred))

3. 实际运行结果

我使用 scikit-learn 1.8.0 和固定随机种子 42 进行了测试：

可以看到，剪枝后叶节点减少约 37%，测试准确率反而有所提高。

原因在于，部分复杂分支可能学习了训练集中的局部噪声。删除这些分支后，模型在当前测试集上的泛化效果有所改善。

不过，剪枝并不保证准确率一定提高。实际开发中应使用验证集选择剪枝参数，并保留独立测试集评价最终效果。

三、案例二：PyTorch 神经网络权重剪枝

1. 业务背景

第二个案例使用手写数字识别任务。

scikit-learn 的 Digits 数据集包含 1,797 张手写数字图像，每张图像的尺寸为 8×8。

我们使用一个简单的两层全连接神经网络进行分类：

64维图像特征
      |
    Linear
      |
   64个神经元
      |
     ReLU
      |
    Linear
      |
   10个类别

训练完成后，使用 PyTorch 的 torch.nn.utils.prune 工具对权重进行剪枝。

本案例采用 L1 非结构化剪枝。

其基本方法是根据权重的绝对值大小，选择绝对值较小的权重，将其置零。

2. Python 完整实践

import numpy as np
import torch

from torch import nn
from torch.nn.utils import prune

from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

# 1. 加载数据
X, y = load_digits(return_X_y=True)

X_train, X_test, y_train, y_test = train_test_split(
    X / 16.0, y,
    test_size=0.2,
    stratify=y,
    random_state=42
)

torch.manual_seed(42)

# 2. 创建神经网络
model = nn.Sequential(
    nn.Linear(64, 64),
    nn.ReLU(),
    nn.Linear(64, 10)
)

xt = torch.tensor(X_train, dtype=torch.float32)
yt = torch.tensor(y_train, dtype=torch.long)
x_test = torch.tensor(X_test, dtype=torch.float32)

optimizer = torch.optim.Adam(
    model.parameters(), lr=0.01
)

# 3. 训练函数
def train(epochs):

    model.train()

    for epoch in range(epochs):

        for idx in torch.randperm(len(xt)).split(128):

            output = model(xt[idx])

            loss = nn.functional.cross_entropy(
                output, yt[idx]
            )

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

# 4. 评价函数
def evaluate():

    model.eval()

    with torch.no_grad():
        pred = model(x_test).argmax(dim=1)

    return np.mean(pred.numpy() == y_test)

# 5. 原始模型训练
train(20)
print("原始准确率:", evaluate())

# 6. 剪掉40%的权重
for layer in [model[0], model[2]]:
    prune.l1_unstructured(
        layer,
        name="weight",
        amount=0.4
    )

print("剪枝后准确率:", evaluate())

# 7. 微调模型
train(5)
print("微调后准确率:", evaluate())

# 8. 固化剪枝结果
for layer in [model[0], model[2]]:
    prune.remove(layer, "weight")

3. 实际运行结果

使用 PyTorch 2.10.0 CPU 环境进行测试：

可以看到，将约 40% 的权重置零后，模型仍保持了接近原始水平的识别准确率。

这一结果说明，该模型中存在一定程度的权重冗余。

但要注意，PyTorch 的 l1_unstructured 会将权重置零，原有矩阵形状保持不变。即使调用 prune.remove() 固化剪枝结果，普通稠密计算也不会自动获得相应比例的加速。

如果希望真正减少计算量和模型存储空间，还需要考虑结构化剪枝、稀疏计算格式或针对部署硬件的优化。

四、实际项目应该怎样选择剪枝方法？

对于神经网络，常见的实践流程是：

训练原始模型 → 评价参数重要性 → 逐步剪枝 → 微调 → 独立测试 → 部署性能测试。

剪枝比例不宜一开始就设得过高。

例如，可以分别尝试 10%、30%、50%、70%，观察模型准确率、稀疏度和实际推理速度的变化。

在设备资源有限的环境中，还可以将剪枝与量化（Quantization）结合，进一步降低模型部署成本。

五、总结

剪枝是一种通过减少冗余结构来优化机器学习模型的技术。

本文两个案例体现了剪枝的不同价值：

决策树剪枝通过减少不必要的分支，降低模型复杂度，并可能改善泛化能力。

神经网络剪枝通过删除或置零贡献较小的权重，提高模型稀疏度，在适当的实现与部署条件下为压缩和加速提供可能。

实际使用时，最重要的是同时评价三个指标：模型预测性能、结构或参数规模，以及真实推理成本。

剪枝的目标是以尽可能小的性能损失，获得更加简洁、高效、容易部署的模型。对于资源受限的设备、边缘计算和大型神经网络，这项技术具有重要的工程应用价值。