yolov5实战之模型剪枝_yolov5模型剪枝

前沿

  在上一篇yolov5的博客中，我们用yolov5训练了一个二维码检测器，可以用来检测图像中是否有二维码，后续可以接一个二维码解码器，就可以解码出二维码的信息了（后续可以聊聊）。这篇博客再讲讲另一个方面：模型轻量化，具体的是轻量化中的模型剪枝。

为什么要做轻量化

  我们训练的模型不仅仅会用在GPU这种算力高的硬件上，也有可能用在嵌入式CPU或者NPU上，这类硬件算力往往较低，尽管在这些设备上运行模型时，我们可以将模型量化为int8，可以大大降低计算量，但有时候只靠这一方式也是不够的。比较直观能想到的提升模型运行速度的方式是裁剪模型，比如减少通道数或模型的深度，这种方式是以牺牲模型精度为代价的。这就促使我们寻找更好的模型轻量化方法，剪枝就是一种使用比较广泛的模型轻量化方法。

什么是剪枝

  模型剪枝（Model Pruning）是一种通过减少神经网络模型中的冗余参数和连接来优化模型的方法。它旨在减小模型的大小、内存占用和计算复杂度，同时尽可能地保持模型的性能。

  模型剪枝的基本思想是通过识别和删除对模型性能影响较小的参数或连接，以达到模型精简和优化的目的。方法包括剪枝后的参数微调、重新训练和微调整体网络结构等。直观的理解就是像下图这样。

  模型剪枝可以在不显著损失模型性能的情况下，大幅度减少模型的参数量和计算量，从而提高模型的部署效率和推理速度。它特别适用于嵌入式设备、移动设备和边缘计算等资源受限的场景，以及需要部署在较小存储空间或带宽受限环境中的应用。
本文选择的模型剪枝方法：Learning Efficient Convolutional Networks through Network Slimming
源代码：https://github.com/foolwood/pytorch-slimming
这个方法基于的想法是通过稀疏化训练，通过BN层的参数，自动得到权重较小通道，去掉这些通道，从而达到模型裁剪的目的。

稀疏化训练

  如上文述，为了达到剪枝的目的，我们要使用稀疏化训练，以使得让模型权重更紧凑，能够去掉一些权重较小的通道，达到模型裁剪的目的。
为了进行稀疏化训练，引入一个稀疏化稀疏参数，这个参数越大，模型越稀疏，能够裁剪的比例越大，需要在实际中调整，参数过大，模型性能可能会下降较多，参数过小，能够裁剪的比例又会过小。
  为了进行稀疏化训练，首先汇总模型的所有BN层：

    if opt.sl > 0:
        print("Sparse Learning Model!")
        print("===> Sparse learning rate is ", hyp['sl'])
        prunable_modules = []
        prunable_module_type = (nn.BatchNorm2d, )
        for i, m in enumerate(model.modules()):
            if isinstance(m, prunable_module_type):