SS-Model【6】：U2-Net

系列文章目录

U-Net语义分割系列（一）：
SS-Model【5】：U-Net
U-Net语义分割系列（二）：
SS-Model【6】：U2-Net

---

---

前言

U2Net 是基于 Unet 提出的一种新的网络结构，同样基于 encode-decode，作者参考 FPN 和 UNet，在此基础之上提出了一种新模块 RSU（ReSidual U-blocks)。 经过测试，对于分割物体前、背景取得了惊人的效果。同样具有较好的实时性，经过测试在 P100 上前向时间仅为18ms（56fps）。

原论文链接：
U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection

---

1. Abstract & Introduction

1.1. Abstract

在本文中，作者设计了一个简单而强大的深度网络架构 -- U-Net，用于显著性物体检测（SOD）。我们的 U-Net 的结构是一个两级嵌套的 U 型结构。该设计有以下优点：

• 由于我们提出的 ReSidual U-block（RSU）中不同大小的感受野的混合，它能够从不同的尺度捕捉更多的上下文信息；
• 由于这些 RSU 块中使用的池化操作，它增加了整个架构的深度，而没有明显增加计算成本。

这种架构使我们能够从头开始训练一个深度网络，而不使用来自图像分类任务的 backbone。

1.2. Introduction

显著性目标检测（SOD）指的是将图像中最显著的物体进行分割，广泛用于视觉追踪和图像分割任务中。

显著性目标检测任务与语义分割任务非常相似，只不过显著性目标检测任务是二分类任务，它的任务是将图片中最吸引人的目标或区域分割出来（偏主观），故只有前景和背景两类。

下图是从 DUTS-TR 数据集中随便挑的几张图片，第一行是原图，第二行是对应的GT，其中白色区域对应前景（最吸引人的目标或区域）黑色区域对应背景：

目前现状：

大多数的 SOD 网络有一个共性，就是注重利用现有的主干提取深层特征，比如 Alexnet、VGG、ResNet、ResNeXt、DenseNet 等。这些主干最终都是为图像分类任务而设计的，他们提取的特征代表语义，而不是局部细节和全局对比信息，但这对显著性检测至关重要。

他们需要在 ImageNet 数据集上进行预训练，如果目标数据与 ImageNet 具有不同的分布，则会比较低效。

SOD 模型中存在的问题：

• 网络结构复杂，这是由于在现有主干网络上添加特征聚合模块，从这些模型中提取多层显著性特征
• 现有主干网络通常通过牺牲高分辨率的特征映射来实现更深层次的体系结构

U2Net 网络的解决方法：

• 该网络是一个两层嵌套的U型结构，没有使用图像分类的预训练主干模型，可以从头训练
• 新的体系结构允许网络更深入、获得高分辨率，而不会显著增加内存和计算成本。
  • 在底层，设计了一个新的 RSU，能够在不降低特征映射分辨率的情况下提取级内多尺度特征
  • 在顶层，有一个类似于 U-Net 的结构，每一 stage 由一个 RSU 块填充

---

2. Network Architecture

2.1. Residual U-blocks

2.1.1. Comparison

对于显著目标检测和其他分割任务来说，局部和全局上下文信息都非常重要。在现代 CNN 设计中，如 VGG、ResNet、DenseNet 等，1×1 或 3×3 的小型卷积滤波器是最常用的特征提取元件，因为它们需要较少的存储空间并且计算效率高。

• ( a ) - ( c ) 显示了具有最小感受野的现有卷积块。由于 1x1 或 3x3 滤波器的感受野太小而无法捕捉全局信息，因此底层的输出特征图只包含局部特征。为了在高分辨率浅层特征图上获取更多的全局信息，最直接的想法是扩大感受野。
• ( d ) 显示了一个 inception like block，试图通过使用空洞卷积来扩大感受野以提取局部特征和非局部特征。然而，在原始分辨率的输入特征图上进行多次扩展卷积（尤其是初始阶段）需要大量的计算和内存资源

2.1.2. RSU - 7

受 U-Net 的启发，提出了一种新的 RSU 来捕捉阶段内多尺度特征。下图显示了 RSU-L(Cin, M, Cout) 结构， 其中 L 是编码器层数，$C_{in}$，$C_{out}$ 表示输入和输出通道，M 表示 RSU 内部层中的通道数。

标出每个输出特征图的 shape 以及 concat 的位置后的重绘结构图如下所示：

RSU 主要由三个部分组成：

• 输入卷积层：它将输入特征图 $H\times W\times C_{in}$ 转成一个具有 Cout 通道数的中间图 $F_1(x)$，这是提取局部特征的普通卷积层
• 以中间特征图 $F_1(x)$ 为输入，学习提取和编码多尺度上下文信息 $U(F_1(x))$
  • U 表示如上图所示的 U-Net
    • L 越大，RSU 越深，池化操作越多，即更大的感受野以及更丰富的局部和全局特征
    • 配置此参数可以从具有任意空间分辨率的输入特征图中提取多尺度特征。从梯度降采样特征图中提取多尺度特征，并通过逐步上采样、合并和卷积等方法将其编码到高分辨率特征图中
  • 这一过程减少了大尺度直接上采样造成的细节损失
• 通过求和： $F_1(x)+U(F_1(x))$ 融合局部特征和多尺度特征

2.1.3. RSU - 4F

与上述 RSU-L 结构不同的是，在 RSU-4F 中并没有进行下采样或上采样，而是将采样层全部替换成了膨胀卷积。作者解释说，到 En_5 时，特征图的分辨率已经很低了，如果接着下采样会丢失很多上下文信息，所以在 RSU-4F 中就不再进行下采样了。这意味着 RSU-4F 的所有中间特征图都与其输入的特征图具有相同的分辨率

上图中带参数 d 的卷积层全部是膨胀卷积，d 为膨胀系数

2.1.4. RSU - 7 & residual block

上图表示残差块与 RSU 对比，主要设计区别在于：

• RSU 用 U-Net 代替了普通的单流卷积，并用一个权重层构成的局部特征代替了原始特征
  • $H_{RSU}(x)=U(F_1(x))+F_1(x)$，其中 U 代表前面图中所示的多层 U 型结构
• 这种设计变化使网络能够直接从每个残差块的多个尺度中提取特征
• 同时，U结构的计算开销很小，因为大多数操作都应用于下采样的特征映射

2.2. Model Structure

作者提出了一种嵌套的 UNet 结构，即在一个大的 UNet 当中，嵌套了一堆小的 UNet。理论上，可以将指数 n 设为任意正整数，实现单级或多级嵌套 U 型结构。但是，嵌套层太多的体系结构过于复杂，无法在实际中实现和应用。

本文所讲解的模型，就是指数为 2 时的一个 2 层嵌套结构，如下图所示。它的顶层是一个由11 stages（图中的立方体）组成的大 U 型结构，每一个 stage 由一个配置良好的 RSU 填充。因此，嵌套的U结构可以更有效的提取 stage 内的多尺度特征和聚集阶段的多层次特征。

在 Encoder 阶段，每通过一个 block 后都会下采样 2 倍（maxpool），在 Decoder 阶段，每通过一个 block 前都会上采样 2 倍（bilinear）

如上图所示，U2-Net网络由三部分构成：

• 六级编码器
  • 编码器 En_1、En_2、En_3 和 En_4 阶段中，使用的是 2.1.2 小节中讲解的 RSU-L 结构，分别为 RSU-7、RSU-6、RSU-5 和 RSU-4
    • 区别在于，后一个编码器比前一个编码器上一步上采样过程和一步下采样过程，下采样倍数也会相应地减少
  • 编码器 En_5、En_6 阶段中，使用的是 2.1.2 小节中讲解的 RSU-4F 结构
• 五级解码器
  • 解码阶段具有与对称编码阶段相似的结构
  • 每个解码器阶段将来自前一级的上采样特征映射和来自其对称编码器阶段的特征映射的级联作为输入
• 与解码器和最后一级编码器相连的显著图融合模型（saliency map fusion module）
  • 通过该模块将不同尺度的 saliency map 进行融合并得到最终预测概率图
  • U2-Net 网络首先通过大小为 $3\times 3$、kernel = 1 的卷积（输出的特征层的 channel 为 1）和 Sigmoid 函数从 En_6，De_5，De_4，De_3，De_2 和 De_1 生成 6 个输出显著概率图 $S_{side}^{(6)}$，$S_{side}^{(5)}$，$S_{side}^{(4)}$，$S_{side}^{(3)}$，$S_{side}^{(2)}$，$S_{side}^{(1)}$
  • 将输出的显著图的逻辑图（卷积输出、Sigmoid 函数之前）向上采样至与输入图像大小一致，并通过级联操作（concat）相融合，然后通过 $1\times 1$ 卷积层和一个 Sigmoid 函数，以生成最终的显著性概率映射图 $S_{fuse}$

2.3. Loss Function

在 U2Net 中损失计算公式如下所示：

$$L=\sum _{m=l}^M{w}_{side}^{(m)}{l}_{side}^{(m)}+w_{fuse}{l}_{fuse}$$

参数含义：

• l 代表二值交叉熵损失
• w 代表每个损失的权重

该损失函数可以看成两部分：

• 一部分是上述提到的 Sup1、Sup2、Sup3、Sup4、Sup5 和 Sup6 与 GT 之间的损失
  • 注意，在计算损失前需要将 Sup1、Sup2、Sup3、Sup4、Sup5 和 Sup6 通过Sigmoid 激活函数得到对应的概率图)
• 另一部分是最终融合得到的概率图与 GT 之间的损失

标准二进制交叉熵计算公式如下所示：

$$l=-\sum _{(r,c)}^{(H,W)}[P_{G(r,c)}log{P}_{S(r,c)}+(1-P_{G(r,c)})log(1-P_{S(r,c)})]$$

参数含义：

• (r, c) 为像素坐标
• (H, W) 为图像大小
• $P_{G(r,c)}$ 表示 GT 像素值
• $P_{S(r,c)}$ 表示预测的显著概率图

训练过程试图最小化整个损失。测试过程中，我们选择最后融合结果 $l_{fuse}$ ​
作为最终的显著性图

2.4. Evaluation indicators

• PR curve
• weighted F-measure
• S-measure
• relax boundary F-measure
• F-measure
  $$F_{\beta }=\frac{(1+{\beta }^2\times Precision\times Recall)}{{\beta }^2\times Precision+Recall}$$
  • 参数含义
    • precision 即为预测为正确的数据中，真实值为正确的比例
    • recall 即为在所有的真实值为正确的数据中，有多少能预测正确
    • 在源代码中 ${\beta }^2=3$
  • 计算得到的范围为：$(0,1)$
  • 数值越大，表示网络的分割效果越好
  • 针对不同的概率阈值会得到一组不同的 precision 和 recall，所以计算得到的 $F_{\beta }$ 是一组数据，最终的指标，是这组数据中最大的值
• MAE ( Mean Absolute Error )
  $$MAE=\frac{1}{H\times W}\sum _{r=1}^H\sum _{c=1}^W|P(r,c)-G(r,c)|$$
  • 参数含义
    • $P(r,c)$ 表示预测概率图
    • $G(r,c)$ 表示 GTboxes
  • 计算得到的范围为：$(0,1)$
  • 数值越小，表示网络的分割效果越好

---

3. DUTS

DUTS 数据集包含了 10553 张训练图片，5019 张测试图片。其中所有的训练图片采集自ImageNet DET 训练 / 验证集，而所有的测试图片采集自 ImageNet DET 测试集以及 SUN 数据集

数据集结构如下所示：

├── DUTS-TR
│      ├── DUTS-TR-Image: 该文件夹存放所有训练集的图片
│      └── DUTS-TR-Mask: 该文件夹存放对应训练图片的GT标签（Mask蒙板形式）
│
└── DUTS-TE
       ├── DUTS-TE-Image: 该文件夹存放所有测试（验证）集的图片
       └── DUTS-TE-Mask: 该文件夹存放对应测试（验证）图片的GT标签（Mask蒙板形式）

---

总结

U2-Net 网络的设计允许具有丰富多尺度特性和相对较低的计算和内存成本de1深层架构。该结构只建立在 RSU 块上，没有使用任何特性分类的预训练主干网络，因此是灵活的，可适应不同的工作环境，性能损失很小。

本文中，使用不同的滤波器配置提供两种情况下的U2-Net

• 普通版本的 U2-Net（176.3MB）
• 较小版本的U2-Net†（4.7MB）

博文参考
视频资料