pytorch_YOLOX剪枝【附代码】

爱吃肉的鹏

已于 2023-04-06 15:13:22 修改

阅读量4.4k

点赞数 22

文章标签：深度学习目标检测 pytorch 剪枝计算机视觉

于 2022-07-04 21:31:18 首次发布

本文链接：https://blog.csdn.net/z240626191s/article/details/125608886

版权

功能更新：2023.4.6新增heatmap可视化

环境

pytorch 1.7

loguru 0.5.3

NVIDIA 1650 4G

intel i5-9th

torch-pruning 0.2.7

安装包

pip install torch_pruning

Note:本项目是在b站up主Bubbliiiing和原YOLOX官方代码进行了整合。

1.添加了feature可视化功能

2.训练中可开启EMA功能

3.网络剪枝(支持s,m,l,x)

3.1支持单个卷积剪枝

3.2支持网络层剪枝

4.剪枝后微调训练

5.Conv与BN层的融合推理加速

6.保存log信息

数据集格式：采用voc数据集格式

feature视化

在tools/Net_Vision.py为可视化代码实现。可以通过在网络层导入NV函数，实现通道可视化。

eg:

features = [out_features[f] for f in self.in_features]
[x2, x1, x0] = features  # shape is (batch_size,channels,w,h)
NV(x2)

python demo.py --predict --heatmap --phi s --heatmap_save_path model_data/heatmap_vision.png

注意：这里的热力图和上面的可视化不一样，上面的可视化是可以对YOLOX中任意卷积层进行可视化，增强网络的可解释能力。而这里的热力图是可以看网络最终的输出。

网络剪枝

参考论文：Pruning Filters for Efficient ConvNets

导入剪枝工具

import torch_pruning as tp

如果需要看yolov4的，可以看：YOLOv4剪枝【附代码】_爱吃肉的鹏的博客-CSDN博客_yolov4剪枝

采用通道剪枝，而不是权重剪枝。

在剪枝之前需要通过tools/prunmodel.py save_whole_model(weights_path, num_classes) 函数将模型的权重和结构都保存下来。

weights_path:权重路径

num_classes：自己类别数量

model = YOLOX(num_classes, 's')  # 这里需要根据自己的类数量修改  s指yolox-s

支持对某个卷积的剪枝：调用Conv_pruning(whole_model_weights):

pruning_idxs = strategy(v, amount=0.4)  # 0.4是剪枝率 根据需要自己修改，数越大剪的越多

对于单独一个卷积的剪枝，需要修改两个地方值，这里的卷积层需要打印模型获得，不要自己盲目瞎猜：

if k == 'backbone.backbone.dark2.0.conv.weight'

pruning_plan = DG.get_pruning_plan((model.backbone.backbone.dark2)[0].conv, tp.prune_conv, idxs=pruning_idxs)

支持网络层的剪枝：调用layer_pruning(whole_model_weights):

included_layers = list((model.backbone.backbone.dark2.modules()))  # 针对某层剪枝

Note:剪枝成功以后，会打印模型的参数变化量！如果没有打印，说明你剪的不对，好好检查一下！

剪枝以后的log日志文件会保存在logs文件下

剪枝后的微调训练

将train.py中的pruned_train设置为True.

False为正常训练，然后自己修改batch_size。

注意修改model_path和classes_path，不然会报错！

剪枝前的网络输入大小和微调训练以及推理时的大小必须一致！

训练自己的数据集

如果你有用过Bubbliiiing up主的代码，你将很快就能上手。数据集采用的是VOC的形式

VOCdevkit/
`-- VOC2007
    |-- Annotations  (存放xml标签文件)
    |-- ImageSets
    |   `-- Main
    `-- JPEGImages (存放图像)

在model_data中新建一个new_classes.txt，里面写入自己的类。运行voc_annotation.py，会在当前目录生成2007_train.txt文件和2007_val.txt文件。(可以检查一下里面有没有生成成功)

在train.py中，将classes_path修改为model_data/new_classes.txt【等预测的时候，也是需要在yolo.py中修改这里】

然后根据需要修改其他超参即可训练，训练权重会保存在logs文件中(默认保存权值，不含网络结构)

预测

参数说明：下面终端的输入都是可选的

--predict:预测模式

--pruned:开启剪枝预测或训练

--image:图像检测

--video:开始视频检测

--video_path:视频路径

--camid:摄像头id 默认0

--fps:FPS测试

--dir_predict:对一个文件夹下图像进行预测

--phi:可以选择s,m,l,x等

--input_shape:网络输入大小，默认640

--confidence:置信度阈值

--nms_iou:iou阈值

--num_classes:类别数量，默认80

--fuse:是否开启卷积层和BN层融合加速，默认False

终端输入：

# 图像预测
python demo.py --predict --image

# 视频预测
python demo.py --predict --video --camdi 0

# fps测试
python demo.py --predict --fps

默认预测都为yolox_s，如果要指定其他网络，输入：(需要注意的是在yolo.py修改权重路径，如果是自己数据集，还需要修改classes_path)

# 使用yolox_l预测
python demo.py --predict --image --phi l

Conv与BN层的融合推理加速

其他命令可以搭配使用，比如采用conv和bn融合的方式进行推理

python demo.py --predict --image --fuse

通过测试发现FPS提升了3帧/s左右(我的GPU是1650)

日志文件的保存

本项目采用loguru工具捕获日志，检测和训练中的一些日志记录会自动记录，保存在logs文件下，一个log文件的大小我设置的上限大小为1 MB，如果超过该范围，会自动生成一个新的.log文件，可以自己修改这个值，或者修改日志保存时间(以免保存了太多的日志)。如果不想要这个功能，可以找到相应的位置注释掉即可。

这里只是帮助大家造轮子，用尽可能简单的代码实现一些功能，不用大家再去看复杂的工程代码，最终的效果需要自己耐心调试，慢慢“炼丹”！