如何对已训练完成的目标检测模型执行剪枝操作？

前置准备 #

• 安装适配TF1.14的模型优化工具包，对应版本命令：pip install tensorflow-model-optimization==0.2.1，更高版本的优化工具包不再支持TF1.x系列。
• 提前保留好已训练模型的checkpoint文件、训练用的pipeline.config配置文件，剪枝微调阶段会直接复用原有模型结构配置。

剪枝操作步骤 #

步骤1：调整训练配置添加剪枝规则 #

在原有的pipeline.config的train_config字段下新增剪枝参数，不需要修改原有模型定义代码：

• 设置剪枝起始步长：建议设置为原模型总训练步长的10%，避免刚启动就剪枝破坏已经收敛的权重特征
• 设置剪枝结束步长：建议设置为原模型总训练步长的50%~70%，预留足够的步长给剪枝后的模型恢复精度
• 设置目标稀疏度：首次尝试建议设为0.5_{0.7，即剪掉50%}70%的无效权重，平衡压缩比和精度损失

注意：剪枝逻辑默认作用于所有卷积层，如果你要保留骨干网络的底层特征权重，可以单独配置排除指定层的剪枝规则。

步骤2：加载预训练权重启动剪枝微调 #

剪枝不是直接删除已有权重，而是先通过稀疏训练让不重要的权重逐步趋近于0，使用OD API自带的训练脚本启动即可，命令示例：

python model_main.py \
--pipeline_config_path=你的pipeline.config文件路径 \
--train_dir=剪枝训练输出目录 \
--fine_tune_checkpoint=你已训练完成的checkpoint文件前缀 \
--num_train_steps=剪枝后总训练步长

训练过程中可通过TensorBoard查看pruning/*类指标，确认稀疏度在按照预设规则逐步上升。

步骤3：导出轻量化推理模型 #

剪枝微调完成后，导出推理模型时需要移除训练阶段添加的剪枝辅助节点，调用OD API自带的导出脚本即可，命令示例：

python export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path 你的pipeline.config文件路径 \
--trained_checkpoint_prefix 剪枝训练完成的checkpoint文件前缀 \
--output_directory 剪枝模型输出目录 \
--strip_pruning_vars true

导出后的模型会自动对值为0的权重做稀疏存储，模型体积、参数量都会按照你设置的稀疏度对应下降。

后续优化建议 #

• 如果剪枝后精度下降超出预期，可以降低目标稀疏度，或者延长微调训练步长，给模型更多恢复精度的空间。
• 剪枝完成后可配合INT8量化操作，使用OD API自带的量化导出参数即可，模型体积还能再降低75%左右，推理速度也会有明显提升。

内容的提问来源于stack exchange，提问作者Sohail