4、权重剪枝实战：使用 tfmot.sparsity.keras 对全连接网络进行权重剪枝

好，咱们直接进入正题。

上一章我们聊了剪枝的理论，说白了就是「把不重要的权重砍掉」。但理论归理论，真正上手写代码，你会发现坑还挺多的。今天我就带你手撸一个完整的权重剪枝实战，用的是 TensorFlow Model Optimization Toolkit 里的 tfmot.sparsity.keras。

我个人习惯，学一个新工具，先跑通一个最简单的例子。所以咱们就拿一个全连接网络开刀，数据集用 MNIST。别嫌简单，麻雀虽小五脏俱全，剪枝的核心流程都在里面了。

4.1 环境准备与依赖安装

首先，确保你装了正确的版本。我踩过这个坑——版本不匹配，API 直接报错。

pip install tensorflow==2.12.0
pip install tensorflow-model-optimization==0.7.5

嗯，这里要注意：tfmot 的版本要和 TensorFlow 对应。我建议你直接用我上面写的版本组合，省心。

4.2 构建一个基准全连接网络

先搭一个没剪枝的模型，作为对比基准。这样你才能看到剪枝后到底损失了多少精度。

import tensorflow as tf
import tensorflow_model_optimization as tfmot
import numpy as np

# 加载 MNIST
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# 构建模型
def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(512, activation='relu'),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    return model

model = build_model()
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5, validation_split=0.1)
baseline_acc = model.evaluate(x_test, y_test, verbose=0)[1]
print(f'基准模型准确率: {baseline_acc:.4f}')

跑完你应该能看到 97% 以上的准确率。这个结果记下来，后面要跟剪枝后的模型做对比。

4.3 应用权重剪枝：核心三步骤

剪枝的 API 其实不复杂，但有几个关键参数你得理解透。我曾经在项目里因为一个参数设错了，剪完的模型直接废了，找 bug 找了半天。

4.3.1 用 tfmot.sparsity.keras.prune_low_magnitude 包装模型

这个函数就是剪枝的核心。它会自动给每一层加上一个「剪枝遮罩」，训练的时候只更新没被剪掉的权重。

pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.30,
        final_sparsity=0.80,
        begin_step=0,
        end_step=1000,
        frequency=100
    )
}

model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(
    build_model(), **pruning_params
)

这里我解释一下这几个参数：

  • initial_sparsity：初始剪枝比例。我习惯设 30%，让模型慢慢适应。
  • final_sparsity：最终剪枝比例。80% 是个比较激进的值，但全连接网络扛得住。
  • begin_step / end_step：从第几步开始剪，到第几步结束。注意这里的「步」是 batch 的步数，不是 epoch。
  • frequency：每隔多少步更新一次遮罩。设 100 比较稳妥。
注意：end_step 不要设太大。如果你训练 5 个 epoch，每个 epoch 有 600 个 batch，那总步数就是 3000。end_step 设 1000 意味着前 1/3 的训练时间完成剪枝，后面 2/3 用来微调。这个比例我个人觉得比较合理。

4.3.2 编译并训练剪枝模型

这里有个小坑：剪枝模型在训练时需要额外的回调函数 UpdatePruningStep。没有它，遮罩不会更新。

model_for_pruning.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

callbacks = [
    tfmot.sparsity.keras.UpdatePruningStep(),
    tfmot.sparsity.keras.PruningSummaries(log_dir='./logs')
]

model_for_pruning.fit(
    x_train, y_train,
    epochs=5,
    validation_split=0.1,
    callbacks=callbacks
)

PruningSummaries 这个回调是可选的，但我强烈建议加上。它会在 TensorBoard 里记录剪枝比例的变化，方便你观察剪枝过程是否平稳。

4.3.3 评估剪枝后的模型

pruned_acc = model_for_pruning.evaluate(x_test, y_test, verbose=0)[1]
print(f'剪枝后模型准确率: {pruned_acc:.4f}')
print(f'准确率下降: {baseline_acc - pruned_acc:.4f}')

我跑出来的结果，80% 剪枝率下准确率大概下降 0.5%~1%。如果你发现下降超过 2%，说明剪枝太猛了，可以降低 final_sparsity 试试。

4.4 导出剪枝模型：去掉遮罩，真正变小

训练完的模型其实还带着遮罩，权重并没有真的变成 0。要真正获得模型大小的收益，需要做一步「剥离」操作。

model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)
model_for_export.save('./pruned_mnist_model.h5')

这一步很多人会忘。我曾经就犯过这个错——训练完直接保存，结果模型大小一点没变，白忙活一场。

对比一下文件大小：

模型 文件大小 准确率
基准模型 约 8.2 MB 97.8%
剪枝后（80%稀疏） 约 2.1 MB 97.2%

你看，大小压缩了将近 4 倍，准确率只掉了 0.6%。这个 trade-off 我觉得非常划算。

4.5 避坑指南与个人经验

经验一：剪枝率不是越高越好。我试过 90% 的剪枝率，准确率直接掉到 85% 以下。对于全连接网络，80% 是个比较安全的阈值。
经验二：如果你发现剪枝后模型不收敛，试试把学习率调低一点。剪枝过程本身会引入噪声，学习率太大容易震荡。
经验三：不要对 BatchNormalization 层做剪枝。BN 层的权重是统计量，剪了会出大问题。tfmot 默认会跳过 BN 层，但如果你自定义了层，记得手动排除。

4.6 小结

好了，到这里你已经掌握了权重剪枝的完整流程。说白了就三步：

  1. 用 prune_low_magnitude 包装模型
  2. 加上 UpdatePruningStep 回调训练
  3. 用 strip_pruning 导出真正压缩后的模型

下一章我们会把这个剪枝后的模型转换成 TensorRT 格式，在 GPU 上跑出更极致的推理速度。到时候你会发现，剪枝 + TensorRT 的组合拳，效果真的绝了。

← 上一章 📖 返回目录 下一章 →