【项目实战】YOLOv3/v4模型加速

  

---

稀疏训练

去除.weights中的epoch信息

稀疏训练

剪枝

对best.pt进行剪枝

对last.pt进行剪枝

微调

.pt文件转.weights文件

---

 　　之前的博客分享过YOLOv3/v4算法训练自己数据的方式，也介绍过模型加速的一般流程，以及本篇模型加速使用的通道剪枝的方式。没做笔记的同学可以回看之后再来阅读本篇噢！

　　此处再来回顾一下模型加速的一般流程。

　　因此，此处省略基础模型训练的环节。

　　本项目实践全仰仗github：https://github.com/tanluren/yolov3-channel-and-layer-pruning，github中ReadMe写的比较清楚，各位同学可以参考ReadMe，也可以参考本篇！

稀疏训练

　　稀疏训练建立在基础模型训练的基础上，也就是YOLOv3/v4原型的训练。（建议使用YOLOv3， YOLOv4在稀疏训练中可能会存在一些问题。可以直接剪枝！）

去除.weights中的epoch信息

　　将训练后的模型放到github项目对应的文件夹中，由于Darknet训练得到的.weights文件是带有epoch记录的，因此，需要先通过转换，将.weights中的epoch信息去掉（这里去掉的含义相当于置零），便于进行稀疏训练。

　　如果基于.weights(带有epoch的)的情况下，需要在训练时将迭代次数设为.weights带有的epoch + 稀疏训练的代数。如果低于.weights带有的epoch，训练时模型会直接保存权重并停止。

　　本篇使用将epoch去掉的方式，将epoch去掉的方法在models.py的convert方法，该方法需要传入两个参数，其一是cfg文件，其二是权重文件。

python -c "from models import *; convert('cfg/yolov3-voc.cfg', 'yolov3-voc.weights')"

　　执行上述代码，会生成convert.pt权重文件(pytorch版的)。

　　之后，便可以使用该权重文件进入稀疏训练的环节。

稀疏训练

　　github上给出了多种剪枝的方式，也给出了相应的调用方法。本篇使用--prune 1适用其他剪枝策略的方式。执行下述代码，将对应的配置文件名称，训练文件名称，权重文件名称替换成自己的；batch-size和device的设置根据自己电脑的配置进行修改，其余参数可根据个人喜好进行修改。如果是yolov4，改成yolov4对应的文件即可。

python train.py --cfg cfg/yolov3-voc.cfg --data data/voc.data --weights yolov3-voc_uav.pt --batch-size 20 --epochs 480 -sr --s 0.005 --prune 1 --device 0

　　代码中每隔10代保存一次权重文件，训练代数设置的比较大的话会十分占硬盘空间的。因此，本篇将保存代数设置为80。

　　训练完成后，会在weights文件夹中生成对应的权重文件，注意区分保存best.pt和last.pt，极容易在后续训练过程中被替换掉。

　　同时，可以打开runs文件夹，对相应的events文件使用TensorBoard观察训练中的情况：

tensorboard --logdir ./

　　系数训练的情况对剪枝会造成比较大的影响，因此，本篇在使用时设置了480 epoch。这种影响体现在相同的剪枝力度会得到不同的参数量。

剪枝

　　该过程是建立在稀疏训练之后，本篇选用slim_prune.py的剪枝方式，使用的是通道剪枝，指令如下。

python slim_prune.py --cfg cfg/yolov3-voc.cfg --data data/voc.data --weights weights/last.pt --global_percent 0.85 --layer_keep 0.01

　　执行上述指令，会在cfg文件中生成相应的cfg文件，weights文件夹中生成相应的.weights文件。本篇记录了yolov3剪枝力度从0.45-0.85的情况，如下图所示，可以看出不同剪枝力度剪枝后，mAP，参数量，推理时间的对应关系。

　　这里有个比较有意思的事情，就是使用稀疏训练生成的best.pt / last.pt分别进行不同力度的剪枝，会出现不一样的效果：

对best.pt进行剪枝

　　对best.pt进行剪枝力度从0.50-0.85进行尝试，结果如下图所示，可以看出，使用best.pt进行剪枝，mAP随着剪枝力度的增大而增大，呈现一定规律性。但剪枝后mAP始终低于原模型。

对last.pt进行剪枝

　　对last.pt进行剪枝力度从0.45-0.70进行尝试。结果如下图所示，可以看出，使用last.pt进行剪枝，剪枝力度仿佛对mAP没有什么影响，掉点并不明显。并且存在剪枝后mAP高于原模型的情况。

微调

　　微调意义在于对剪枝之后的模型恢复精度。也可以在微调的模型上加大数据集，使得剪枝后的模型泛化能力可以具有与原模型比肩的能力。

python train.py --cfg cfg/prune_0.85_keep_0.01_yolov3-voc.cfg --data data/voc.data --weights weights/prune_0.85_keep_0.01_last.weights --epochs 400 --batch-size 20 --device 0

　　微调后，保存的模型又是.pt文件，会生成和稀疏训练一样的best.pt和last.pt，很有可能会替代之前的文件，如果之前的文件放置在了weights文件目录下。

.pt文件转.weights文件

　　为了便于将生成的权重文件和配置文件放回Darknet中进行测试，需要将微调生成的.pt权重文件转换为.weights文件。

　　由于剪枝过程是使用.pt文件为输入，生成的是.weights文件，个人通过对slim_prune.py代码的阅读，从中抠出了.pt文件转.weights文件的代码，不过仍需要基于该github相关的文件。注意，不是单纯执行下述代码就可以实现.pt文件向.weights文件，需要基于github项目。

# -*- coding: utf-8 -*-
# @Time    : 2020/7/7 上午9:22
# @Author  : monologuesmw   
# @Email   : monologuesmw@163.com
# @File    : pt2weights.py
# @Software: PyCharm

import torch

from models import * 
from utils.prune_utils import *        # 基于github项目中的两个文件

import argparse
import numpy as np

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
    # parser.add_argument('--data', type=str, default='data/coco.data', help='*.data file path')
    parser.add_argument('--weights', type=str, default='weights/last.pt', help='sparse model weights')
    # parser.add_argument('--global_percent', type=float, default=0.8, help='global channel prune percent')
    # parser.add_argument('--layer_keep', type=float, default=0.01, help='channel keep percent per layer')
    parser.add_argument('--img_size', type=int, default=416, help='inference size (pixels)')
    opt = parser.parse_args()

    img_size = opt.img_size
    cfg_path = opt.cfg
    pt_path = opt.weights
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = Darknet(opt.cfg, (img_size, img_size)).to(device)
    assert pt_path
    model.load_state_dict(torch.load(opt.weights, map_location=device)['model'])

    compact_model_name = "my_test.weights"
    save_weights(model, path=compact_model_name)
    print("success save weights")

　　执行下述指令便可以得到相应的my_test.weights文件。

python pt2weights.py --cfg cfg/prune_0.85_keep_0.01_yolov3-voc.cfg --weights weights/best.pt

  

　　（训练、测试数据部分来源于Kaggle）