本文针对yolov5n的模型进行结构化剪枝、微调,在文章最后会放上程序链接。
剪枝方法:基于BN层系数gamma剪枝。
前言:yolov5n模型是非常轻量级的网络模型,但针对嵌入式端依然比较大;此时,可以通过添加L1正则来约束BN层系数,使得系数稀疏化,通过稀疏训练后,裁剪掉稀疏很小的层,对应激活也很小,所以对后面的影响非常小,反复迭代这个过程,可以获得很compact的模型。
原理:
在一个卷积-BN-激活模块中,BN层可以实现通道的缩放。如下:
BN层的具体操作有两部分:
在归一化后会进行线性变换,那么当系数gamma很小时候,对应的激活(Zout)会相应很小。这些响应很小的输出可以裁剪掉,这样就实现了bn层的通道剪枝。
通过在loss函数中添加gamma的L1正则约束,可以实现gamma的稀疏化。
上面损失函数L右边第一项是原始的损失函数,第二项是约束,其中g(s) = |s|,λ是正则系数,根据数据集调整。
实际训练的时候,就是在优化L最小,依据梯度下降算法:
′=∑′+∑′()=∑′+∑||′=∑′+∑∗()
所以只需要在BP传播时候,在BN层权重乘以权重的符号函数输出和系数即可,对应添加如下代码:
# Backward
loss.backward()
# scaler.scale(loss).backward()
# # ============================= sparsity training ========================== #
srtmp = opt.sr*(1 - 0.9*epoch/epochs)
if opt.st:
ignore_bn_list = []
for k, m in model.named_modules():
if isinstance(m, Bottleneck):
if m.add:
ignore_bn_list.append(k.rsplit(".", 2)[0] + ".cv1.bn")
ignore_bn_list.append(k + '.cv1.bn')
ignore_bn_list.append(k + '.cv2.bn')
if isinstance(m, nn.BatchNorm2d) and (k not in ignore_bn_list):
m.weight.grad.data.add_(srtmp * torch.sign(m.weight.data)) # L1
m.bias.grad.data.add_(opt.sr*10 * torch.sign(m.bias.data)) # L1
# # ============================= sparsity training ========================== #
optimizer.step()
# scaler.step(optimizer) # optimizer.step
# scaler.update()
optimizer.zero_grad()
这里并未对所有BN层gamma进行约束,这里对C3结构中的Bottleneck结构中有shortcut的层不进行剪枝,主要是为了保持tensor维度可以加:
稀疏训练后,将BN层前的卷积层对应通道的卷积核裁剪掉,将BN层后对应的特征图裁剪掉。
具体步骤:
1、首先,下载yolov5-6.0(其他版本应该也支持,目前只尝试了6.0/6.1)的源码,下载yolov5n.pt文件,配置训练自己数据集,得到一个模型。
2、然后,进行稀疏化训练;
train_sparity.py
在train.py中添加如下代码:
# Backward
# scaler.scale(loss).backward()
loss.backward()
# # ============================= sparsity training ========================== #
srtmp = opt.sr*(1 - 0.9*epoch/epochs)
if opt.st:
ignore_bn_list = []
for k, m in model.named_modules():
if isinstance(m, Bottleneck):
if m.add:
ignore_bn_list.append(k.rsplit(".", 2)[0] + ".cv1.bn")
ignore_bn_list.append(k + '.cv1.bn')
ignore_bn_list.append(k + '.cv2.bn')
if isinstance(m, nn.BatchNorm2d) and (k not in ignore_bn_list):
m.weight.grad.data.add_(srtmp * torch.sign(m.weight.data)) # L1
m.bias.grad.data.add_(opt.sr*10 * torch.sign(m.bias.data)) # L1
设置好如下的参数,加载刚才训练的模型,即可稀疏化训练:
此时训练容易报错:
ModuleNotFoundError: No module named ‘utils.loggers.wandb’
将下面47行注释掉即可。
稀疏训练前后map值变化如下所示:
稀疏训练前的训练:
稀疏训练50轮:
效果有所增加。
3、训练完成后进行剪枝
prune.py
此时容易出现:TypeError: run() got an unexpected keyword argument 'cfg’错误,
解决方法:在prune.py的 def run 和 def run_prune 中添加模型文件和参数信息即可,切记不可只在 def parse_opt 中添加。
4、微调训练测试
在finetune_prune.py中设置模型和data路径后,运行:
python finetune_pruned.py --weights pruned_model.pt --adam --epochs 50
map0.5值为0.865,相比正常训练时的0.871,精度下降了0.006,但模型大小由3.57兆减小为2.64兆,大小下降了26%,精度下降可忽略不计。
剪枝模型测试:修改detectpruned.py中模型和img路径,运行即可。
程序网盘链接:
链接:https://pan.baidu.com/s/1bGUonuaUPdeoPY-6IXOGPQ?pwd=1234
提取码:1234