基于显著性目标检测的非特定类别图像分割实现以及部署过程

一、前景主体分割要求

• 模型大小：不超过100M
• 算法性能指标：在1080 TI GPU上处理一张图片的时间不超过5秒
• 效果指标：主体边缘细节清晰，MIOU越高越好
• 网络模型工程化：离线可进行验证的包含全部依赖的C++例程

二、实现方案

使用显著性目标检测(SOD)模型U2Net，U2Net提供了源码以及两个模型，一个167.3M，一个4.7M。

至此，将问题转化为将优化U2Net模型大小，以下为实现方案：

• 降低模型精度，使用float16进行计算（失败，预测不正确，得到全黑的图像）
• 将普通的卷积层改为分组卷积，分为2组
• 将普通的卷积层改为深度可分离卷积

三、训练

3.1 数据集

• 训练集
  将DUTS-TR（5019张图像）、DUTS-TE（10553张图片）作为训练集（总共15572张图像）。
• 验证集
  在DUT-OMRON（5168张）、ECSSD（1000张）、HKU-IS（4447张）、MSRA10K（10000张）、MSRA-B（5000张））中（共25615张）随机取1000张作为验证集。

3.2 训练细节

在两张1080 TI上训练，批次大小(batch size)设置为16，每迭代2000次，保存一次模型，
迭代200000次，结束训练（即保存有100个模型），其他参数使用原U2Net的参数。

训练时使用以下两种参数初始化方式：

• 加载修改后U2Net的模型（167.3M）做为初始化参数、
  1. 分组卷积（将卷积层权重相邻的通道，两两切分作为一组，取平均值，得到一个通道，处理权重的所有通道，并将处理后的参数级联到一起。）【模型大小86M】
  2. 深度可分离卷积（将depthwise conv的权重设置为原来普通卷积层(out_channels, in_channels, kernel_h, kernel_width)的第一维取平均得到，再第一维与第二位做转置，得到(in_channels, 1, kernel_h, kernel_width)的张量作为初始化参数；depthwise conv的bias，设置为size为(in_channels)的全零张量，pointwise conv的权重，将原来普通卷积层的张量的第三第四维求平均得到(out_channels, in_channels, 1, 1)大小的张量，pointwise conv的bias，使用原来普通卷积层的bias。）【模型大小21M】
• 直接默认的初始化方式，进行训练。（尝试之后，决定放弃这种训练方式）

加载修改后的模型初始化参数，可以加快损失下降速度（而且不只是快一点点。）如下图为分组卷积在两种不同训练方式下的损失变化。

默认初始化：训练十几轮（每轮2000次迭代）从0.48降到0.4左右

修改后的模型初始化参数：训练三十几轮（每轮2000次迭代）从0.44左右降到0.06左右，十几轮就已经降到了0.1左右

四、部署方案

将pytorch下训练得到的pth模型，转为onnx模型，使用OpenCV中的dnn模块加载进行推理（其中对OpenCV进行重新编译，加入CUDA加速处理相关依赖，时间从1800ms左右，降到了120ms左右，在i5-7400 CPU @ 3.00GHz + 1050Ti的主机上）。

五、完整源码

下载链接https://download.csdn.net/download/DeepLearning_/87378107