对MODNet 主干网络 MobileNetV2的剪枝探索

1 引言

1.1 MODNet 原理

1.2 MODNet 模型分析

2 MobileNetV2 剪枝

2.1 剪枝过程

2.2 剪枝结果

2.2.1 网络结构

2.2.2 推理时延

2.3 实验结论

3 模型嵌入

3.1 模型保存与加载

法一：保存整个模型

法二：仅保存模型的参数

小试牛刀

小结

3.2 修改 MobileNetV2 中的 block

阶段一

阶段二

总结

3.2 MobileNetv2 参数嵌入 v1

3.2.1 遇到问题

3.2.2 解决方案

3.2.3 探索过程

阶段一

阶段二

阶段三

3.2.4 结论

3.3 MobileNetV2 参数嵌入 v2

3.3.1 开展思路✨

3.3.2 探索过程

阶段一

ERROR1

ERROR 2

阶段二

阶段三

模拟推理

ERROR 1

ERROR 2

阶段四

ERROR 1

ERROR 2

实验结果

1 引言

1.1 MODNet 原理

基于卷积神经网络擅长处理单一任务的特性，MODNet 将人像抠图分解成了三个相关的子任务：语义估计(Semantic Estimation)、细节预测(Detail Prediction)、语义-细节融合(Semantic-Detail Fusion)，分别对应着模型结构中 Low-Resolution Branch(LR-Branch)、High-Resolution Branch(HR-Branch)、Fusion Branch(F-Branch) 三个分支。其中，LR-Branch 用于人像整体轮廓的预测，HR-Branch 用于前景人像与背景过渡区域的细节预测， Fusion Branch 负责将前两部分预测得到的轮廓与边缘细节融合。语义估计部分是独立的分支，依赖于细节预测、语义-细节融合两个分支。

MODNet论文地址：https://arxiv.org/pdf/2011.11961.pdf

1.2 MODNet 模型分析

作为轻量化模型，MODNet 参数量仅6.45MB，但作为底层视觉任务，MODNet输入、输出分辨率一致，计算复杂度较高。笔者利用 NNI 分析了 MODNet 三个分支的参数量与计算量：

MODNet分支	参数量(MB)	计算量(GFLOPs)
LR-Branch	6.16	5.03
HR-Branch	0.24	10.58
F-Branch	0.05	2.71

从计算量来看，MODNet 三个分支 LR-Branch、HR-Branch、F-Branch 分别占整个网络的27.7%、58.2%、14.1%。HR-Branch 作为高分辨率分支，相比其它两个分支的计算量较高，是剪枝重点关注的部分。

在实际剪枝时，由于 MODNet 内部包含了三个分支，直接剪枝较为复杂，因此对结构进行拆分。

接下来，是笔者对主干网络 MobileNetV2 部分的探索实验。

2 MobileNetV2 剪枝

2.1 剪枝过程

确定待剪枝的对象：

Cov2d layer
BatchNorm2d
Linear

✨剪枝策略：基于L1范数结构化剪枝

model = MobileNetV2(3)

config_list = [{'sparsity': 0.8, 'op_types': ['Conv2d']}]
pruner = L1NormPruner(model, config_list)
_, masks = pruner.compress()
pruner._unwrap_model()
ModelSpeedup(model, torch.rand(1, 3, 512, 512), masks).speedup_model()

2.2 剪枝结果

参数量：3.5M → 309.76k

复杂度：1.67 GMac → 41.54 MMac

从参数量和计算复杂度来看，对MobileNetV2的剪枝可以明显压缩模型，下面具体来看网络结构、推理时延与精度。

2.2.1 网络结构

分别展示三组，A组、B组、C组分别代表MobileNetV2第0~15层、第16~32层、第33~52层。

A组剪枝前：

A组剪枝后：

B组剪枝前：

B组剪枝后：

C组剪枝前：

C组剪枝后：

可以发现，不论是靠近模型输入端，还是输出端，卷积层都有大幅度的裁剪。

2.2.2 推理时延

使用PyTorch框架，分别在GPU与CPU上测试，结果如下。

GPU上

剪枝前	剪枝后
1	0.9339	0.8370
2	0.9144	0.7980
3	0.9530	0.8759

CPU：

	剪枝前	剪枝后
1	0.3529	0.0609
2	0.3906	0.0549
3	0.3600	0.0520

第一次观察剪枝前、后的推理时延时，发现差距并不大。正在此时，笔者联想到前一秒的数据类型不一的报错：由于 dummy_input 转到了 CUDA 上，而 model 还是 CPU上的 model，导致无法 run。于是笔者把 model 转到了cuda上。也正是在这个时候，笔者想到了推理硬件的问题：如果数据量本身就不大的时候，在 CUDA 上运行速度不一定比 CPU 快。数据如果要利用 GPU 计算，就需要从内存转移到显存上，数据传输具有很大开销，对于小规模数据来说，传输时间已经超过了 CPU 直接计算的时间。因此，规模小无法体现出 GPU 的优势，而大规模神经网络采用 GPU 加速意义较大，推理同样如此。

通过对MobileNetV2剪枝，我们可以发现：剪枝前、后的模型在 CPU 上的推理状况已经有了明显的差距，推理时延的降低达到了预期！！！

仔细看推理，dummy_input 初始化时涉及了batch，上述实验都是基于batch=1的情况。因此，接下来对剪枝前、后的模型，在不同的batch、在不同的硬件上推理进行对比。

CPU：

batch	剪枝前	剪枝后
1	0.30	0.07
2	0.64	0.07
4	1.14	0.18
8	2.22	0.25
16	4.33	0.39
32	11.76	0.78
64	66.45	1.49

GPU：

batch	剪枝前	剪枝后
1	0.81	0.77
2	0.90	0.79
4	0.95	0.76
8	xxx	0.83
16	xxx	0.91
32	xxx	0.97
64	xxx	xxx

2.3 实验结论

回首LeNet，不论是pth格式的PyTorch模拟推理还是实际推理，以及 ONNX、OpenVINO 推理，剪枝前、后的模型在 CPU 上的推理速度差异不明显就可以解释了，同时也验证了先前的猜想：LeNet 本身属于小规模网络，结构简单，推理速度已经很快了。因此，不论如何压缩模型，在batch较小的时候都无法有效降低推理时延。

另外，笔者总结了下列结论🎉：

当网络规模并非很大时，在CPU上推理比GPU更适合；（规模的阈值没有明确界限，需要尝试）

当 batch 逐渐增大时，剪枝后的模型推理速度占有绝对优势，且大模型优于小模型；

若要进行实时推理，保证较低的推理时延，batch 可以考虑设置为1；

相较于小模型，大模型剪枝在实际推理中更有意义；（LeNet剪枝前后batch为1时延不变，而mobilenetv2差异较大）

不仅仅是模型大小，数据本身也会影响推理速度；（推理小数据负担较小，所以在做模型压缩时，是否可以考虑前处理，例如图像压缩）

3 模型嵌入

3.1 模型保存与加载

PyTorch 模型的保存与读取方式有两种，接下来还是以 LeNet 为例，结构定义如下：

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 5 * 5, 120) 
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

法一：保存整个模型

model = LeNet()
torch.save(model,PATH)

# 直接加载，得到网络结构：
model = torch.load(PATH)

💥注意：参数可以通过 model.named_parameters() 获取。

法二：仅保存模型的参数

torch.save(model.state_dict(),PATH)

# 直接加载，得到权重参数：
model = torch.load(PATH)

若要将参数加载到网络中，需要获取网络结构，也就是初始化model：

model = LeNet()
model.load_state_dict(torch.load('model2.pth'))
print(model)

💥注意：参数和网络结构应当相匹配，否则参数无法传入。

这里做个示范：将第二个全连接层的输出通道以及第三个全连接层的输入通道都改为94，继续load，则：

因此，在对网络结构裁剪后，利用 state_dict 读取参数前需要相应修改网络结构，可以参考裁剪后的模型通道数以及卷积核的变化情况。

小试🐂🔪

以 LeNet 为例，由于之前都是直接保存整个结构，因此没有涉及到结构的修改问题，接下来只保存模型参数。

剪枝后的网络结构为：

加载剪枝后的模型：

model = LeNet()
path = 'new_model.pth'
model.load_state_dict(torch.load(path))
print(model)

所以，如果直接加载则会产生参数与结构不匹配的错误，因此修改如下：

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 2, 3)
        self.conv2 = nn.Conv2d(2, 4, 3)
        self.fc1 = nn.Linear(4 * 5 * 5, 120) 
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

小结

模型保存时只保存参数，相对保存整个模型而言，不需要在加载时与结构的定义在同一脚本中；此外仅保存模型的参数占用较少的内存，且节省时间，所以一般情况下只保存参数的做法更常用。

3.2 修改 MobileNetV2 中的 block

阶段一

修改前、后对比：

# 修改前
    interverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [expansion, 24, 2, 2],
            [expansion, 32, 3, 2],
            [expansion, 64, 4, 2],
            [expansion, 96, 3, 1],
            [expansion, 160, 3, 2],
            [expansion, 320, 1, 1],
        ]
    
# 修改后：
        interverted_residual_setting = [
            # t, c, n, s
            [1, 4, 1, 1],
            [expansion, 8, 2, 2],
            [expansion, 18, 3, 2],
            [expansion, 39, 4, 2],
            [expansion, 45, 3, 1],
            [expansion, 77, 3, 2],
            [expansion, 64, 1, 1],
        ]

然而，由于裁剪后通道数没有保持8的整数倍，加载失败也在预料之中。

🎉解决方案：

利用_make_divisible对NNI剪枝源码进行修改；❌
修改裁剪比例✅

初步探索：

①将根据稀疏比例设定位0.2进行裁剪，保留80%的成分，但又由于浮点数问题的取整问题也无法保证是8的倍数。

②再次观察通道数，只要将第二个残差块的out_channels（24）保留，剩下的裁剪50%即可保证。但通过实验发现，NNI的API在根据预先设定的稀疏度裁剪时存在误差，在对于MobileNetv2这样的倒置残差块结构来说，无法达到理想的效果。

③通过查看源码，发现了剪枝时的mode选项：当设置为“dependency_aware”时，剪枝器将强制具有依赖关系的卷积层剪枝相同的通道。因此，一方面可以保证准确的稀疏度；另一方面，加速模块可以更好地从剪枝模型中获得速度优势。

剪枝配置如下：

config_list = [{'sparsity': 0.5,
                'op_types': ['Conv2d']},
               {'exclude': True,
                'op_names': ['features.0.0','features.2.conv.6']
                }]

这里加入了first layer的conv，如果不排除，也会导致channel不是8的倍数。

这一次达到了预期效果，但加在模型参数时依然提示do not match。

这里思考了很久，再一看模型，发现实际上笔者一直在关注channel是否为8的倍数，而忽略了interverted_residual_block中的扩展因子，每次扩展都会将低维乘6。因此，block内部的维度扩展在剪枝过后不再满足要求，如下图：

阶段二

🎉解决方案：

修改block层的expansion系数；
排除每一个block“低维--->高维”时的第一层进行剪枝；

首先，将低维扩展后的第一个高维conv进行保留测试，剪枝后的shape如下：

按照这样的思路，将其余block的第一个扩展层进行排除，配置如下：

config_list = [{'sparsity': 0.5,
                'op_types': ['Conv2d']},
               {'exclude': True,
                'op_names': ['features.0.0', 'features.2.conv.6', 'features.2.conv.0', 'features.3.conv.0',
                             'features.4.conv.0', 'features.5.conv.0', 'features.6.conv.0', 'features.7.conv.0',
                             'features.8.conv.0', 'features.9.conv.0', 'features.10.conv.0', 'features.11.conv.0',
                             'features.12.conv.0', 'features.13.conv.0', 'features.14.conv.0', 'features.15.conv.0',
                             'features.16.conv.0', 'features.17.conv.0']
                }]

整体上达到了理想效果，仅有一个block层不符合expansion，随之修改对应block层的expansion以及each block setting。

由于该层会重复执行两次，在将expansion修改为3以后，第二次执行就不再符合条件。这里根据通道变化情况针对执行次数单独控制，如下：

        interverted_residual_setting = [
            # t, c, n, s
            [1, 8, 1, 1],
            [6, 24, 2, 2],  # 1. 8-->24  2. 24-->24,expansion=3(modify)
            [6, 16, 3, 2],  # 1. 24-->16(modify)  2. 16-->96
            [expansion, 32, 4, 2],
            [expansion, 48, 3, 1],
            [expansion, 80, 3, 2],
            [expansion, 160, 1, 1],
        ]

成功加载：

mobilenet = MobileNetV2(3)
mobilenet.load_state_dict(torch.load('modified_mobilenet_3.pth'))
print(mobilenet)

接下来，针对MODNet预训练模型中的MobileNetV2部分进行裁剪。

第一次剪枝下来有个别层的通道有误，同时，个别通道没有满足预设的expansion，比如：

小结

考虑到计算机处理器单元架构的问题，通道数保持8的倍数能够较快计算，因此，剪枝也应当遵循这个原则；

考虑到每一个block内部的扩展因子expansion，因此，排除这些层中的第一个连接层以满足倍数关系；

对仍不满足要求的层特殊处理，对相同blcok不同执行次数时的expansion控制；

3.2 MobileNetv2 参数嵌入 v1

3.2.1 遇到问题

将MobileNetV2结构嵌入 MODNet 并不困难，难点在于基于MODNet预训练模型参数裁剪后的嵌入、进而推理。

在将裁剪后的MobileNetV2模型结构直接嵌入MODNet，开始训练时，提示通道数不匹配的error，原因在于MobileNetV2只是MODNet中的一部分。

因此，针对这样的问题，笔者考虑了下列四种方案。

3.2.2 解决方案

方案一：保留 MobileNetV2 模型的 input 以及 output channel，重新裁剪；

方案二：修改 MobileNetV2 前、后模型的关联通道数；

方案三：将裁剪后的预训练模型参数加载到修改后的 MobileNetV2，将其余参数正常加载到其余三个子模块，最后合并；

方案四：直接对 MODNet 进行裁剪；

3.2.3 探索过程

阶段一

易由于已经对 MODNet 预训练参数中的 MobileNetV2 参数部分进行了裁剪，同时也修改了 MobileNetV2 结构，在单独对这一子模块进行测试时可以成功嵌入。基于这样的情况，首先采用方案三。

参数合并具体作法：通过对象访问类的成员，即lr_branch、hr_branch、f_branch，在获取其结构后，从MODNet预训练参数中load对应部分的参数，再结合已有的backbone，分别替换原来MODNet中的四个子模块。

替换 backbone 部分：

mobilenet2 = my_mobilenetv2_v2.MobileNetV2(3)
mobilenet2.load_state_dict(torch.load('my_mobilenetv2_2.pth'), strict=False)

modnet.MODNet().backbone = mobilenet2
print(modnet.MODNet().backbone)

实际上，该方法是无效的。通过比对替换前后的MODNet中的backbone参数，可以发现并非理想的数值。

阶段二

因此，笔者开始寻找有效的参数替换方法，参考这位博主。

接下来，在LeNet上进行实验。

按照博客上的代码跑下来出现了一个bug，即缺乏 state_dict 属性，原因在于只保存了模型参数，而非整个模型。只有完整的模型才拥有状态字典。如下：

由于本次剪枝是针对MODNet预训练模型参数，因此，这里按照博客的思路重新调整语句，如下：

model0 = LeNet()
model1 = LeNet()

# 同一结构，不同参数
path0 = 'new_model.pth'
path1 = 'new_model1.pth'

params0 = torch.load(path0)
params1 = torch.load(path1)

new_param = {}
for name, params in model0.named_parameters():
    new_param[name] = params0[name] + params1[name]

torch.save(new_param, 'new.pth')
model0.load_state_dict(torch.load('new.pth'))
print(list(model0.named_parameters()))

params0：

params1：

new_params：(+)

结果是将两个同一结构、不同参数的模型，进行求和运算，得到新模型，或者说，实现了参数修改。

💥注意：模型结构不变，只有参数在变化。

基于此，回到MODNet，继续实现。

此时又意识到了一个问题：MODNet中的子结构发生了变化，也就是说，不仅仅是参数的替换。

这里有一种方案：

1、根据剪枝后的情况修改网络结构，得到新结构，并保存随机参数；

2、将新结构中不含backbone部分的参数用预训练模型中的参数替换，backbone部分使用裁剪后的参数；

3、保存，得到新结构、新参数

在实验过程中遇到了一个error：加载预训练模型以后，每一次打印的参数不一致。

阶段三

采用方案四，即直接对MODNet模型裁剪。

通过打印输出可发现，MODNet总共有72个卷积层，2个全连接层。由于两个全连接层权重以及输入尺寸较大，因此占有较多的参数，同时，与第二个卷积层相连接的卷积层也有着巨大的参数量。如下：

而这一卷积层的参数量将近整个MODNet的1/2，因此，对该层的裁剪不容忽视。

首先排除剪枝时无法被NNI支持的算子：

算子排除有问题，偶然间看到了新版NNI，安装以后对官方代码进行测试：

虽然NNI2.8 的确可以成功 run 带有 expand_as 的算子，但从问题本身来看，该算子没有进行处理，即只是官方在源码中对该算子进行了排除，并非如 aten_relu 算子般检测到以后使用 torch.relu() 进行操作。

3.2.4 结论

虽然 MobileNetV2 只是 MODNet 的一部分，在裁剪时也是针对这部分进行了操作，但在修改 MobileNetV2 结构以后，从整个 MODNet 结构来看，与 MobileNetV2 相邻的结构也相应产生了变化，也正表明：刚开始的方案三，针对MobileNetV2以及其余子模块分别处理的不合理性，而方案一和二也就没有了必要。

参数修改的有效办法并非简单进行结构赋值，而是通过 state_dict 根据键值对替换；

在利用多块 GPU 并行计算时，需要利用 DataParallel，在保存参数时需要使用 model.module.state_dict()，而非像往常一样的 model.state_dict()。虽然后者不会报错，但需要通过指定 load 时 strict 为 false，才可以将参数名匹配。当然，如果采用前者保存，一方面不需要在 load 时指定 strict，同时也能保证加载到模型的参数名保持一致。

aten:expand_as 算子在新版的 NNI 源码中被排除；

对 MODNet 子模块剪枝较容易，但难以将参数嵌入；

3.3 MobileNetV2 参数嵌入 v2

获取网络块，在将 MODNet 参数加载到结构后，通过对象访问获取 MobileNetv2，参数可以保持一致。

model = modnet.MODNet(backbone_pretrained=False)
pretrained_ckpt = 'modnet_photographic_portrait_matting.ckpt'
model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(pretrained_ckpt).items()})
model = model.backbone  # 获取主干网络：mobilenet V2

3.3.1 开展思路✨

按照先前的 config 设定进行剪枝，并对 MobileNetV2 网络进行修改；
将修改后的 MobileNet V2 嵌入 MODNet，得到新的 MODNet，保存一份带有随机初始化参数的模型；
按照上述获取 MobileNet V2 参数的方式，获取其余三个子模块的参数，保存；
读取2中保存的模型，用3得到的权重参数进行替换；
保存新模型；

3.3.2 探索过程

阶段一

ERROR1

分析原因：网络结构中backbone外面包裹着model，而模型参数保存时缺少，因此参数不匹配。

解决方案：可以通过指定strict = False，当然从根源上来看，修改模型保存时的指令（类似多卡训练保存！）

# 将 torch.save(model.state_dict(), PATH)替换为下列语句：

torch.save(model.model.state_dict(), PATH)

ERROR 2

分析原因：classifier层的参数缺失，因此保存参数时也就缺少了这一部分的权重参数。

实际上，先前剪枝同样遇到了这样的问题，但在加载时通过指定strict = False，对齐了参数名以及结构。

解决方案：暂时先以同样的方式对齐，解决了该问题。

与此同时，证明了这一结论：如果模型仅保存了部分参数，当加载到网络中时，部分参数可以正常填入，而剩余的部分会随机初始化。

进一步观察网络结构可以发现，MobileNet V2 中包含 classifier 层，但由于该层用于分类任务，因此在 MODNet 使用时将其去除，这也就成了剪枝后 classifier 参数无法对应填入的原因。

解决方案：修改MobileNet V2网络结构，去除 classifier 的定义。

实验结果：

在去除分类层以后，参数与结构一一对应，因此也就不存在上述假设的问题了。

至此，MobileNet V2 的剪枝、结构修改都已完成。

阶段二

先将裁剪后的 MobileNet V2 的参数替换至 MODNet MobileNet V2 的随机初始化参数：

import torch
from src.models.modnet import MODNet

model = MODNet(backbone_pretrained=False)
model_params = torch.load("pruned_backbone_state.pth")

new_params = {}

for name, params in model.named_parameters():  # 随机初始化参数，将整个MODNet中的层都打印了出来
    if 'backbone' in name:  # 将裁剪后的参数成功填入
        new_name = name.replace('backbone.model.', '')
        new_params[name] = model_params[new_name]
    else:
        new_params[name] = params
        
# print(list(model.named_parameters()))
torch.save(new_params, 'new_substitution_model.pth')

替换前后对比如下：

替换前MODNet头部：

MobileNet V2头部：

替换以后MODNet头部：

此时，新模型除了MobileNet V2以外，其他参数都是随机的。

接下来，利用新结构加载该模型参数。

这里的bug让笔者觉得奇怪，在原先查看权值参数时，每一个layer仅仅只有weight与bias，但这里却出现了其他属性。

检查MODNet所有的参数名并检查替换后的state_dict：

model = MODNet(backbone_pretrained=False)
d = {name: param for name, param in model.named_parameters()}
print(d.keys())

print(list(new_params.keys()))

通过比对两者参数，可以发现：数量与参数名都一致。

有一个值得考虑的地方：在参数替换时利用了 named_parameters() 获取模型参数，因此检验两者固然相同。但是，named_parameters() 本身是否真的获取到了模型所有参数呢？

查阅资料可知：named_parameters() 对于获取BN layer来说，只有两个可学习参数alpha、belta，而running_mean以及running_var，作为两个不可学习的统计量只有在model.state_dict()的时候输出。其作用在于作为训练过程对batch的数据统计。

基于此，由于这些参数对实际推理不起作用，直接用规定strict进行加载。

阶段三

模拟推理

import torch
from src.models.modnet import MODNet
from time import time

model = MODNet(backbone_pretrained=False)
model.load_state_dict(torch.load('new_substitution_model.pth'), strict=False)
print(list(model.named_parameters()))

# inference
dummy_input = torch.randn([1, 3, 512, 512]).to('cpu')
t1 = time()
a = model(dummy_input)
t2 = time()
print("Infer time: ", t2 - t1)