pytorch深度学习+目标检测+训练自己数据集+YOLOX+剪枝+轻量化程序源代码Pruning_for_YOLOX.rar
(10.5 MB, 需要: RMB 29 元)
采用pytorch深度学习环境,目标检测为YOLOX,可做剪枝。功能描述:1.可训练自己数据集2.支持图像推理3.支持视频推理4.支持mAP测试5.支持任意层的剪枝6.支持多种模型7.支持微调训练内含readme文件,可快速上手。代码成熟,已有多人测试适用适用人群:研究人员,研究生,大学生,深度学习研究人员。
[TOC]
# 环境
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:
```python
features = [out_features[f] for f in self.in_features]
[x2, x1, x0] = features # shape is(batch_size,channels,w,h)
NV(x2)
```
<p align="center">
<img src="img/特征图.jpg"/>
</p>
# 网络剪枝
参考论文:**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):**
```python
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的形式
。。。。。
。


雷达卡




京公网安备 11010802022788号







