本项目实现功能如下：
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 （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 yolov5n.yaml 1KB共 109 条
- 1
- 2
资源评论
