本项目实现功能如下:
1.训练自己的数据集
2.对任意卷积层进行剪枝
3.剪枝后的训练
4.剪枝后的模型预测
5.利用知识蒸馏对剪枝后模型进行训练
```
环境说明:
gitpython>=3.1.30
matplotlib>=3.3
numpy>=1.18.5
opencv-python>=4.1.1
Pillow>=7.1.2
psutil # system resources
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
thop>=0.1.1 # FLOPs computation
torch>=1.7.0 # see https://pytorch.org/get-started/locally (recommended)
torchvision>=0.8.1
tqdm>=4.64.0
ultralytics>=8.0.100
torch_pruning==0.2.7
pandas>=1.1.4
seaborn>=0.11.0
```
# 1.训练自己的数据集
将自己制作好的数据集放在dataset文件下,目录形式如下:
> dataset
> |-- Annotations
> |-- ImageSets
> |-- images
> |-- labels
**Annotations是存放xml标签文件的**,**images是存放图像的**,**ImageSets存放四个txt文件**【后面运行代码的时候会自动生成】,labels是将xml转txt文件。
> 1.运行makeTXT.py。这将会在ImageSets文件夹下生成 trainval.txt,test.txt,train.txt,val.txt四个文件【如果你打开这些txt文件,里面仅有图像的名字】。
> 2.打开voc_label.py,并修改代码 classes=[""]填入自己的类名,比如你的是训练猫和狗,那么就是classes=["dog","cat"],然后运行该程序。此时会在labels文件下生成对应每个图像的txt文件,形式如下:【最前面的0是类对应的索引,我这里只有一个类,后面的四个数为box的参数,均归一化以后的,分别表示box的左上和右下坐标,等训练的时候会处理成center_x,center_y,w, h】
> ```
> 0 0.4723557692307693 0.5408653846153847 0.34375 0.8990384615384616
> 0 0.8834134615384616 0.5793269230769231 0.21875 0.8221153846153847
> ```
3.在data文件夹下新建一个mydata.yaml文件。内容如下【你也可以把coco.yaml复制过来】。
你只需要修改nc以及names即可,nc是类的数量,names是类的名字。
> ```
> train: ./dataset/train.txt
> val: ./dataset/val.txt
> test: ./dataset/test.txt
>
> # number of classes
> nc: 1
>
> # class names
> names: ['target']
> ```
4.终端输入参数,开始训练。
以yolov5s为例:
```python
python train.py --weights yolov5s.pt --cfg models/yolov5s.yaml --data data/mydata.yaml
```
> from n params module arguments
> 0 -1 1 3520 models.common.Conv [3, 32, 6, 2, 2]
> 1 -1 1 18560 models.common.Conv [32, 64, 3, 2]
> 2 -1 1 18816 models.common.C3 [64, 64, 1]
> 3 -1 1 73984 models.common.Conv [64, 128, 3, 2]
> 4 -1 2 115712 models.common.C3 [128, 128, 2]
> 5 -1 1 295424 models.common.Conv [128, 256, 3, 2]
> 6 -1 3 625152 models.common.C3 [256, 256, 3]
> 7 -1 1 1180672 models.common.Conv [256, 512, 3, 2]
> 8 -1 1 1182720 models.common.C3 [512, 512, 1]
> 9 -1 1 656896 models.common.SPPF [512, 512, 5]
> 10 -1 1 131584 models.common.Conv [512, 256, 1, 1]
> 11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
> 12 [-1, 6] 1 0 models.common.Concat [1]
> 13 -1 1 361984 models.common.C3 [512, 256, 1, False]
> 20 -1 1 296448 models.common.C3 [256, 256, 1, False]
> 21 -1 1 590336 models.common.Conv [256, 256, 3, 2]
> 22 [-1, 10] 1 0 models.common.Concat [1]
> 23 -1 1 1182720 models.common.C3 [512, 512, 1, False]
> 24 [17, 20, 23] 1 16182 models.yolo.Detect [1, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
> Model Summary: 270 layers, 7022326 parameters, 7022326 gradients, 15.8 GFLOPs
> Starting training for 300 epochs...
>
> Epoch gpu_mem box obj cls labels img_size
> 0/299 0.589G 0.0779 0.03841 0 4 640: 6%|████▋ | 23/359 [00:23<04:15, 1.31it/s]
>
>
看到以上信息就开始训练了。
# 2.对任意卷积层进行剪枝
在利用剪枝功能前,需要安装一下剪枝的库。需要安装0.2.7版本,0.2.8有粉丝说有问题。剪枝时的一些log信息会自动保存在logs文件夹下,每个log的大小我设置的为1MB,如果有其他需要大家可以更改。
```python
pip install torch_pruning==0.2.7
```
YOLOv5与我之前写过的剪枝不同,v5在训练保存后的权重本身就保存了完整的model,即用的是torch.save(model,...),而不是torch.save(model.state_dict(),...),因此不需要单独在对网络结构保存一次。
模型剪枝代码在tools/prunmodel.py。你只需要找到这部分代码进行修改:我这里是以剪枝整个backbone的卷积层为例,如果你要剪枝的是其他层**按需修改.included_layers内就是你要剪枝的层。**
```python
"""
这里写要剪枝的层
"""
included_layers = []
for layer in model.model[:10]:
if type(layer) is Conv:
included_layers.append(layer.conv)
elif type(layer) is C3:
included_layers.append(layer.cv1.conv)
included_layers.append(layer.cv2.conv)
included_layers.append(layer.cv3.conv)
elif type(layer) is SPPF:
included_layers.append(layer.cv1.conv)
included_layers.append(layer.cv2.conv)
```
接下来在找到下面这行代码,amount为剪枝率,同样也是按需修改。【这里需要明白的一点,这里的剪枝率仅是对你要剪枝的所有层剪枝这么多,并不是把网络从头到尾全部剪,有些粉丝说我选了一层,剪枝率50%,怎么模型还那么大,没啥变化,这个就是他搞混了,他以为是对整个网络剪枝50%】。
```python
pruning_plan = DG.get_pruning_plan(m, tp.prune_conv, idxs=strategy(m.weight, amount=0.8))
```
接下来调用剪枝函数,传入参数为自己的训练好的权重文件路径。
```python
layer_pruning('../runs/train/exp/weights/best.pt')
```
见到如下形式,就说明剪枝成功了,剪枝以后的权重会保存在model_data下,名字为layer_pruning.pt。
这里需要说明一下,保存的权重文件中不仅包含了网络结构和权值内容,还有优化器的权值,如果仅仅保存网络结构和权值也是可以的,这样pt会更小一点,我这里默认都保存是为了和官方pt格式一致。
```python
-------------
[ <DEP: prune_conv => prune_conv on model.9.cv2.conv (Conv2d(208, 512, kernel_size=(1, 1), stride=(1, 1), bias=False))>, Index=[0, 1, 2, 3, 7, 8, 10, 11, 12, 13, 16, 17, 18, 19, 21, 22, 23, 25, 27, 28, 29, 30, 31, 32, 33, 34, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 53, 54, 56, 57, 58, 59, 60, 61, 62, 63, 65, 67, 69, 70, 71, 72, 73, 74, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 89, 90, 91, 92, 95, 96, 97, 99, 100, 102, 103, 104, 105, 106, 107, 109, 110, 111, 113, 114, 115, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 132, 133, 135, 137, 139, 142, 143, 144, 146, 148, 150, 152, 153, 154, 155, 156, 157, 158, 159, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 173, 174, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 215, 216, 217, 219, 220, 221, 222, 223, 224, 225, 226, 228, 229, 230, 232, 233, 234, 235, 236, 237, 239, 240, 241, 242, 243, 246, 247, 248, 249, 251, 252, 253, 254, 257, 258, 259, 260, 263, 264, 265, 266, 267, 268, 270, 271, 272, 273, 274, 275, 276, 277, 278, 280, 281, 282, 283, 284, 285, 286, 287, 288, 2
没有合适的资源?快使用搜索试试~ 我知道了~
温馨提示
人工智能-项目实践-知识蒸馏-基于对剪枝后的yolov5模型进行知识蒸馏训练python源码.zip 本项目实现功能如下: 1.训练自己的数据集 2.对任意卷积层进行剪枝 3.剪枝后的训练 4.剪枝后的模型预测 5.利用知识蒸馏对剪枝后模型进行训练 1.训练自己的数据集 将自己制作好的数据集放在dataset文件下,目录形式如下: dataset |-- Annotations |-- ImageSets |-- images |-- labels Annotations是存放xml标签文件的,images是存放图像的,ImageSets存放四个txt文件【后面运行代码的时候会自动生成】,labels是将xml转txt文件。 1.运行makeTXT.py。这将会在ImageSets文件夹下生成 trainval.txt,test.txt,train.txt,val.txt四个文件【如果你打开这些txt文件,里面仅有图像的名字】。 2.打开voc_label.py,并修改代码 classes=[""]填入自己的类名,比如你的是训练猫和狗,那么就是classes=["do
资源推荐
资源详情
资源评论
收起资源包目录
人工智能-项目实践-知识蒸馏-基于对剪枝后的yolov5模型进行知识蒸馏训练python源码.zip (109个子文件)
Dockerfile 821B
bus.jpg 476KB
zidane.jpg 165KB
readme.md 22KB
README.md 10KB
README.md 2KB
datasets.py 43KB
train_dil.py 33KB
general.py 33KB
train.py 32KB
wandb_utils.py 25KB
common.py 22KB
tf.py 20KB
plots.py 19KB
val.py 17KB
export.py 16KB
yolo.py 16KB
detect.py 15KB
torch_utils.py 14KB
metrics.py 13KB
augmentations.py 11KB
loss.py 11KB
auto_label.py 10KB
autoanchor.py 7KB
__init__.py 6KB
hubconf.py 6KB
downloads.py 6KB
experimental.py 4KB
activations.py 4KB
prunmodel.py 4KB
printmodel.py 3KB
voc_label.py 2KB
callbacks.py 2KB
gen_wts.py 2KB
resume.py 1KB
restapi.py 1KB
sweep.py 989B
log_dataset.py 891B
testmode.py 867B
makeTXT.py 861B
example_request.py 299B
__init__.py 0B
__init__.py 0B
__init__.py 0B
__init__.py 0B
datasets.cpython-37.pyc 34KB
general.cpython-37.pyc 29KB
common.cpython-37.pyc 24KB
wandb_utils.cpython-37.pyc 18KB
plots.cpython-37.pyc 17KB
yolo.cpython-37.pyc 12KB
torch_utils.cpython-37.pyc 12KB
metrics.cpython-37.pyc 10KB
augmentations.cpython-37.pyc 9KB
loss.cpython-37.pyc 8KB
__init__.cpython-37.pyc 6KB
autoanchor.cpython-37.pyc 6KB
experimental.cpython-37.pyc 5KB
downloads.cpython-37.pyc 4KB
callbacks.cpython-37.pyc 2KB
__init__.cpython-37.pyc 148B
__init__.cpython-37.pyc 135B
__init__.cpython-37.pyc 134B
userdata.sh 1KB
get_coco.sh 900B
mime.sh 780B
get_coco128.sh 615B
download_weights.sh 443B
additional_requirements.txt 105B
Objects365.yaml 7KB
xView.yaml 5KB
VOC.yaml 3KB
anchors.yaml 3KB
VisDrone.yaml 3KB
Argoverse.yaml 3KB
sweep.yaml 2KB
SKU-110K.yaml 2KB
coco.yaml 2KB
yolov5-p7.yaml 2KB
yolov5s.yaml 2KB
GlobalWheat2020.yaml 2KB
yolov5x6.yaml 2KB
yolov5s6.yaml 2KB
yolov5n6.yaml 2KB
yolov5m6.yaml 2KB
yolov5l6.yaml 2KB
coco128.yaml 2KB
hyp.scratch-low.yaml 2KB
hyp.scratch-high.yaml 2KB
hyp.scratch.yaml 2KB
yolov5-p6.yaml 2KB
yolov5-p2.yaml 2KB
yolov3-spp.yaml 2KB
yolov3.yaml 2KB
yolov5-panet.yaml 1KB
yolov5s-ghost.yaml 1KB
yolov5s-transformer.yaml 1KB
yolov5-bifpn.yaml 1KB
yolov5m.yaml 1KB
yolov5x.yaml 1KB
共 109 条
- 1
- 2
资源评论
.whl
- 粉丝: 3823
- 资源: 4648
上传资源 快速赚钱
- 我的内容管理 展开
- 我的资源 快来上传第一个资源
- 我的收益 登录查看自己的收益
- 我的积分 登录查看自己的积分
- 我的C币 登录后查看C币余额
- 我的收藏
- 我的下载
- 下载帮助
最新资源
资源上传下载、课程学习等过程中有任何疑问或建议,欢迎提出宝贵意见哦~我们会及时处理!
点击此处反馈
安全验证
文档复制为VIP权益,开通VIP直接复制
信息提交成功