https://blog.csdn.net/IEEE_FELLOW/article/details/117236025作者只提供了原版的yolov5，不能自定义剪枝层。如果只对网络的头部剪枝，而主干不剪枝，则上述代码无法实现。因此便有了本文。如需代码请评论区联系。
本文以mobilenetv3-yolov5为例
如何将yolov5改为mobilenetv3-yolov5不是本文重点，但这里还是提一下。只需修改common.py yolo.py添加mobilenetv3的定义。 然后使用自己的配置文件即可。
下面给两个参考：
https://blog.csdn.net/weixin_44808161/article/details/125759652
https://blog.csdn.net/m0_58996495/article/details/127998636目录
1 稀疏训练
这里的1458 是主干网络中MobileNet_Block的卷积名称。可以根据网络层数查看 m 的名称，假如主干网络是0-12层（训练时会打印log，显示每层的class），那么逐层查看名称，找到其命名规律。本文忽略了主干网络所有卷积层，所以把所有名称都加入到忽略列表，即其bn参数不参与稀疏化训练。
2.剪枝
修改配置 ：数据集、模型、剪枝比例。
接下来不妨边trace 代码边修改
（1）line 864 ：先测试然后剪枝
（2）忽略列表
下面是忽略列表，忽略的部分不参与剪枝。加入主干网络的所有卷积层。
（3）定义剪枝吼模型的基本配置文件
剪枝完成后会保存每一层的输入通道数和输出通道数，因此需要定义一个类接受这些参数。即使没对主干网络剪枝为了代码规范统一也要定义新的MobileNet_Block。具体的，主要在原来MobileNet_Block的基础上加入input和output参数，有多少卷积层就加多少组。
（4）执行剪枝
得到剪枝后每层通道数
不用修改代码
（5）根据剪枝mask组装模型
进入此方法（函数）
下图是关键之处，进入此方法
进入之后新增代码
详细代码贴在下面
elif m in [conv_bn_hswish]:
named_m_bn = named_m_base + ".bn"
bnc = int(maskbndict[named_m_bn].sum())
c1, c2 = ch[f], bnc
args = [c1, c2, *args[1:]]
layertmp = named_m_bn
if i>0:
from_to_map[layertmp] = fromlayer[f]
fromlayer.append(named_m_bn)
分析下上述代码
conv_bn_hswish
named_m_bn = named_m_base + ".bn"
mask中该层对应的bn层的名称，这个可以通过遍历网络查查。
bnc = int(maskbndict[named_m_bn].sum())
该层的mask中的输出通道是多少
c1, c2 = ch[f], bnc
c1是输入通道数，即为上层网络的输出。c2是输出通道数。
args = [c1, c2, *args[1:]]
*args[1:]表示在（3）中定义的配置文件中除了第一个参数外的其他参数
layertmp = named_m_bn
if i>0:
from_to_map[layertmp] = fromlayer[f]
fromlayer.append(named_m_bn)
将本层网络加入到fromlayer中。from_to_map这个是映射表，表示此层网络来来源于哪里。
MobileNet_BlockPruned （工作量大）
elif m in [MobileNet_BlockPruned]:
convs_num = 0
for key in list(maskbndict.keys()):
if named_m_base.split(".")[-1] == key.split(".")[1]:
convs_num = convs_num + 1
if convs_num == 3:
named_m_cv1_bn = named_m_base + ".conv.1"
named_m_cv2_bn = named_m_base + ".conv.4"
named_m_cv3_bn = named_m_base + ".conv.8"
from_to_map[named_m_cv1_bn] = fromlayer[f]
# from_to_map[named_m_cv2_bn] = fromlayer[f]
fromlayer.append(named_m_cv3_bn)
cv1in = ch[f]
cv1out = int(maskbndict[named_m_cv1_bn].sum())
hiddendim = int(maskbndict[named_m_cv2_bn].sum())
outp = int(maskbndict[named_m_cv3_bn].sum())
args = [cv1in, cv1out, hiddendim, outp, args[-6], args[-5], args[-4], args[-3], args[-2], args[-1]]
c2 = outp
elif convs_num == 2:
named_m_cv1_bn = named_m_base + ".conv.1"
named_m_cv2_bn = named_m_base + ".conv.5"
from_to_map[named_m_cv1_bn] = fromlayer[f]
# from_to_map[named_m_cv2_bn] = fromlayer[f]
fromlayer.append(named_m_cv2_bn)
cv1in = ch[f]
cv1out = cv1in
hiddendim = int(maskbndict[named_m_cv1_bn].sum())
outp = int(maskbndict[named_m_cv2_bn].sum())
args = [cv1in, cv1out, hiddendim, outp, args[-6], args[-5], args[-4], args[-3], args[-2], args[-1]]
c2 = outpfor key in list(maskbndict.keys()):
if named_m_base.split(".")[-1] == key.split(".")[1]:
convs_num = convs_num + 1
MobileNet_BlockPruned有时有两个卷积层，有时有三个卷积层，所以这里要判断下。下面的代码是根据卷积层的个数来装载不同的参数。
（6）复制权重 （工作量大）
返回这里。下面的大循环是复制权重的操作
下文给出for循环的详细代码
# # add by fby
for ((layername, layer), (pruned_layername, pruned_layer)) in zip(model.named_modules(),
pruned_model.named_modules()):
assert layername == pruned_layername
layerindex = -1
if len(layername.split(".")) > 1:
layerindex = int(layername.split(".")[1])
if isinstance(layer, nn.Conv2d) and (layerindex == 0 or layerindex > 11) and not layername.startswith(
"model.26"):
convname = layername[:-4] + "bn"
if convname in from_to_map.keys():
former = from_to_map[convname]
if isinstance(former, str):
out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername[:-4] + "bn"].cpu().numpy())))
in_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[former].cpu().numpy())))
w = layer.weight.data[:, in_idx, :, :].clone()
if len(w.shape) == 3: # remain only 1 channel.
w = w.unsqueeze(1)
w = w[out_idx, :, :, :].clone()
pruned_layer.weight.data = w.clone()
changed_state.append(layername + ".weight")
if isinstance(former, list):
orignin = [modelstate[i + ".weight"].shape[0] for i in former]
formerin = []
for it in range(len(former)):
name = former[it]
tmp = [i for i in range(maskbndict[name].shape[0]) if maskbndict[name][i] == 1]
if it > 0:
tmp = [k + sum(orignin[:it]) for k in tmp]
formerin.extend(tmp)
out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername[:-4] + "bn"].cpu().numpy())))
w = layer.weight.data[out_idx, :, :, :].clone()
pruned_layer.weight.data = w[:, formerin, :, :].clone()
changed_state.append(layername + ".weight")
else:
keyy = layername[:-4] + "bn"
valuee = maskbndict[keyy]
aaa = np.argwhere(np.asarray(valuee.cpu().numpy()))
out_idx = np.squeeze(np.argwhere(np.asarray(maskbndict[layername[:-4] + "bn"].cpu().numpy())))
w = layer.weight.data[out_idx, :, :, :].clone()
assert len(w.shape) == 4
pruned_layer.weight.data = w.clone()
changed_state.append(layername + ".weight")
# fby
elif isinstance(layer, nn.Conv2d) and not layername.startswith("model.26"):
pruned_layer.weight.data = layer.weight.data.clone()layerindex是网络的层数，如果是第0层（主干网络非MobileNet部分），或者是11层以上（非主干网络），不是第26层（detet层），那么是被剪枝过的，要复制对应通道的参数。
这里的26是最后一层的索引，根据自定义网络层数修改
3 微调
简单修改下参数配置即可
扫一扫
7325
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
