首先使用基于yolov5m731.pt进行稀疏训练：
python train_sparsity.py --st --src 0.0001 --srb 0.0001 --weights yolov5m731.pt --adam --epochs 100
src、srb的选择需要根据数据集调整，可以通过观察tensorboard的map，gamma变化直方图等选择。 在run/train/exp*/目录下:
tensorboard --logdir .
然后点击出现的链接观察训练中的各项指标.
训练完成后进行剪枝：
python prune_convbn.py --weights runs/train/exp1/weights/last.pt --conv_percent 0.15 --bn_percent 0.75
裁剪比例conv_percent、bn_percent根据效果调整，可以从小到大试。裁剪完成会保存对应的模型pruned_model.pt。
微调：
python finetune_prune_conv.py --weights pruned_model.pt --adam --epochs 100
| model | pruned | map | mode size |
|---|---|---|---|
| yolov5m731.pt | - | 80.0 | 40.6 M |
| sparity.pt | 0.15\0.75 | 74.6 | 40.6 M |
| pruned.pt | - | 74.6 | 16.0 M |
| fine-tune.pt | - | 70.0 | 8.2 M |
- 浅层尽量少剪,从训练完成后gamma每一层的分布也可以看出来.
- 系数λ的选择需要平衡map和剪枝力度.首先通过train.py训练一个正常情况下的baseline.然后在稀疏训练过程中观察MAP和gamma直方图变化,MAP掉点严重和gamma稀疏过快等情况下,可以适当降低λ.反之如果你想压缩一个尽量小的模型,可以适当调整λ.
- 稀疏训练=>剪枝=>微调 可以反复迭代这个过程多次剪枝.
- 稀疏训练是非常种重要的,也是调参的重点,多观察bn直方图变化,过快或者过慢都不适合,所以需要平衡你的sr, lr等.一般情况下,稀疏训练的结果和正常训练map是比较接近的.
- 剪枝时候多试试不同的ratio,一个基本的准则是每层bn层至少保留一个channel,所以有时候稀疏训练不到位,而ratio设置的很大,会看到remaining channel里面会有0出现,这时候要么设置更小的ratio,要么重新稀疏训练,获得更稀疏的参数.
- 如果想要移植到移动端，可以使用ncnn加速，另外剪枝时控制剩余channel为2^n能有效提升推理速度；GPU可以使用TensorRT加速。