Skip to content

实现对YOLOX的剪枝操作,添加了卷积层和BN层融合推理,添加中间层可视化功能,可实现预测和训练日志保存

Notifications You must be signed in to change notification settings

YINYIPENG-EN/Pruning_for_YOLOX

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

[TOC]

更新:2023.4.6新增headmap可视化

环境

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

将会在model_data下保存热力图。

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

网络剪枝

参考论文:Pruning Filters for Efficient ConvNets

导入剪枝工具

import torch_pruning as tp

如果需要看yolov4的,可以看:https://blog.csdn.net/z240626191s/article/details/124326595

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

在剪枝之前需要通过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_pathclasses_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文件中(默认保存权值,不含网络结构)

计算mAP

先在yolo.py修改model_path为自己需要测试的权重路径

参数说明:

**--map_mode:**指定测试mAP模式,0表示整个mAP计算流程,1表示仅仅获得预测结果,2表示获得数据集的真实框,3表示仅仅获得mAP,4表示coco计算mAP

**--classes_path:**自己数据集的classes txt路径.

--map_vis: 是否开启map 图像可视化

**--voc_path :**VOC数据集根目录

--out_path: 输出路径

**--pruned:**剪枝模式【该模式下指定phi无效

--phi: s,m,l,x

**--input_shape:**输入大小

--confidence: 置信度阈值

**--nms_iou:**NMS阈值

**--fuse:**conv与BN层的融合

Eg:

python get_map.py --map_mode 0 --pruned --classes_path [your classes path]

预测

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

--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,如果要指定其他网络(即指定phi,该模式只在未剪枝时有效),输入:(需要注意的是在yolo.py修改权重路径,如果是预测自己数据集,还需要修改yolo.py中的classes_path和model_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文件,可以自己修改这个值,或者修改日志保存时间(以免保存了太多的日志)。如果不想要这个功能,可以找到相应的位置注释掉即可。

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

权重

链接:https://pan.baidu.com/s/1Jbq8dCv893rZ7RkaANUZgQ 提取码:yypn

About

实现对YOLOX的剪枝操作,添加了卷积层和BN层融合推理,添加中间层可视化功能,可实现预测和训练日志保存

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages