ACC-UNet:面向2020年代的完全卷积UNet模型

论文题目：ACC-UNet: A Completely Convolutional UNet
model for the 2020s
论文地址：https://arxiv.org/pdf/2308.13680

。

1. 背景介绍

语义分割任务的目标是对图像中的每个像素进行分类，这在医学图像分析、自动驾驶和卫星影像处理等领域有着重要的应用。传统的 UNet 模型利用对称的编码器-解码器结构和跳跃连接（skip connections）在这一任务上取得了很好的效果。随着 Transformer 模型的出现，研究者开始探索如何结合 CNN 和 Transformer 的优势，以提升模型的表现能力。

2. 为什么需要改进？

虽然 UNet 和基于 Transformer 的模型（如 Swin-Unet）各有优点，但它们在以下几个方面存在改进空间：

• 卷积网络的局部感受野限制：CNN 依赖固定大小的卷积核进行特征提取，可能会遗漏更大范围的上下文信息。
• 语义差距问题：传统 UNet 的跳跃连接直接将编码器特征传递到解码器，可能导致高层语义信息和低层特征之间的差距，影响分割精度。
• 参数效率和计算效率：基于 Transformer 的模型虽然表现强大，但通常需要大量计算资源和数据才能充分训练。

3. ACC-UNet 的核心创新

为了克服这些问题，ACC-UNet 提出了几项关键创新，试图在保持高效性的同时提升模型性能。

3.1 倒置瓶颈（Inverted Bottleneck）

$$x_1={\text{DConv}}_{3\times 3}\left({\text{PConv}}_{c_{\text{in}}\to c_{\text{inv}}}(x_{\text{in}})\right)$$

• $x_{\text{in}}$ 是输入特征图。
• ${\text{PConv}}_{c_{\text{in}}\to c_{\text{inv}}}$ 表示将输入特征图的通道数从 $c_{\text{in}}$ 通过逐点卷积（pointwise convolution）扩展到 $c_{\text{inv}}$。
• ${\text{DConv}}_{3\times 3}$ 表示在扩展后的特征图上应用 $3\times 3$ 的深度卷积（depthwise convolution）。

什么是倒置瓶颈？

倒置瓶颈 是一种在深度学习中常用的结构，特别是在移动端模型（如 MobileNetV2）中使用。倒置瓶颈的设计理念是：

• 先扩展，后压缩：输入特征图的通道数 $c_{in}$先通过一个逐点卷积（1x1 卷积）扩展到一个较高的维度$c_{inv}=c_{in}$*inv_fctr
• 深度卷积（Depthwise Convolution）：使用 3x3 的深度卷积对扩展后的特征图进行空间维度上的卷积操作，这样操作的计算复杂度较低，因为它没有增加参数量，只是增加了特征图的通道数。
• 再压缩回原始维度：最后，通过另一个逐点卷积将扩展后的特征图的通道数再压缩回到$c_{in}$或$c_{out}$

如果$c_{in}=32$，inv_fctr=3，那么$c_{inv}=32\times 3=96$，输入特征图的通道数从32扩展到96。

3.2 HANC 块（Hierarchical Aggregation of Neighborhood Context）

什么是 HANC 块？

自注意力的简化版本：

• 自注意力会将每个像素与所有其他像素进行比较，这计算量非常大。
• 为了简化这个过程，我们可以通过比较一个像素与其邻域的均值和最大值来近似这种操作。

用均值和最大值来替代自注意力：

• 均值（mean）：计算一个像素周围（邻域）所有像素的平均值。这给出了一个“平均邻域”的概念。
• 最大值（max）：计算一个像素周围所有像素的最大值。这给出了一个“最强邻域”的概念。

层次化（Hierarchical）上下文聚合：

• 考虑图像中的信息通常有层次结构的特点，作者建议不要在一个大的窗口中计算均值和最大值，而是在不同大小的窗口中进行分层计算。
• 比如，我们可以先计算 2x2的窗口，然后是4x4（$2^2\times 2^2$），接着是 8x8 （$2^3\times 2^3$）等等。
• 通过这样逐层增加窗口大小，可以逐渐捕捉到更大的上下文信息，而不需要使用计算复杂度高的大的卷积核。

特征拼接（Concatenation）：

• 为了丰富特征表示，将这些不同尺度下计算得到的均值和最大值与原始特征图进行拼接。
• 拼接操作是沿着通道维度（即深度方向）进行的，这样原始的每个特征图都会增加额外的上下文信息。

公式说明

$$x_2=\left(x_1||{\text{mean}}_{2\times 2}(x_1)||{\text{mean}}_{2^2\times 2^2}(x_1)||\cdots ||{\text{mean}}_{2^{k-1}\times 2^{k-1}}(x_1)||{\text{max}}_{2\times 2}(x_1)||{\text{max}}_{2^2\times 2^2}(x_1)||\cdots ||{\text{max}}_{2^{k-1}\times 2^{k-1}}(x_1)\right)$$

• $x_1$：原始特征图。
• ${\text{mean}}_{2\times 2}(x_1)$：计算每个像素在 $2\times 2$ 邻域内的均值。
• ${\text{max}}_{2\times 2}(x_1)$：计算每个像素在 $2\times 2$ 邻域内的最大值
• 通过这种方式，我们得到了一系列不同尺度的均值和最大值特征图。
• $||$：表示将这些特征图在通道维度上拼接起来。

$$x_3=P\cdot {\text{Conv}}_{c_{\text{inv}}\cdot (2^{k-1})\to c_{\text{in}}}(x_2)+x_{\text{in}}$$

• $x_2$：这是经过多级特征聚合（包括计算均值和最大值）后的特征图，具有更丰富的上下文信息和更多的通道数。此时，$x_2$ 的通道数是 $c_{\text{inv}}\cdot (2^{k-1})$，因为它包含了多级（$2^{k-1}$）聚合的特征。

• $\text{PConv}c\text{inv}\cdot (2^{k-1})\to c_{\text{in}}$：这是一个逐点卷积（Pointwise Convolution）操作，用来改变特征图的通道数。逐点卷积的输入通道数是 $c_{\text{inv}}\cdot (2^{k-1})$，输出通道数是 $c_{\text{in}}$。这一步的目的是将扩展后的特征图 $x_2$ 的通道数重新压缩到与原始输入特征图 $x_{\text{in}}$ 一致。

• $x_{\text{in}}$：这是原始输入特征图，它的通道数是 $c_{\text{in}}$。

• $x_3$：这是最终的输出特征图。它是通过将逐点卷积操作的输出与原始输入特征图 $x_{\text{in}}$ 相加得到的。
  $$x_{\text{out}}=P\cdot {\text{Conv}}_{\text{cin}\to \text{cout}}(x_3)$$
  将特征图$x_3$的通道数从$c_{i}n$转换为$c_{out}$，$c_{out}$为希望得到的特征图的通道数。

举例说明

假设我们有一个输入特征图 $x_1$，其大小为 $3\times 3$，并且它只有一个通道。我们将从简单的 $2\times 2$ 窗口开始计算均值和最大值。这个例子将帮助我们理解计算过程及其用途。

1. 输入特征图 $x_1$

假设特征图 $x_1$ 为：

$$x_1=\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$$

2. 计算 $2\times 2$ 邻域的均值和最大值

首先，我们将计算每个 $2\times 2$ 邻域的均值和最大值。

均值计算：

左上角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}1& 2\\ 4& 5\end{array}\right]$$

$${\text{mean}}_{2\times 2}=\frac{1+2+4+5}{4}=3$$

右上角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}2& 3\\ 5& 6\end{array}\right]$$

$${\text{mean}}_{2\times 2}=\frac{2+3+5+6}{4}=4$$

左下角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}4& 5\\ 7& 8\end{array}\right]$$

$${\text{mean}}_{2\times 2}=\frac{4+5+7+8}{4}=6$$

右下角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}5& 6\\ 8& 9\end{array}\right]$$

$${\text{mean}}_{2\times 2}=\frac{5+6+8+9}{4}=7$$

最大值计算：

左上角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}1& 2\\ 4& 5\end{array}\right]$$

$${\text{max}}_{2\times 2}=\max (1,2,4,5)=5$$

右上角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}2& 3\\ 5& 6\end{array}\right]$$

$${\text{max}}_{2\times 2}=\max (2,3,5,6)=6$$

左下角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}4& 5\\ 7& 8\end{array}\right]$$

$${\text{max}}_{2\times 2}=\max (4,5,7,8)=8$$

右下角 $2\times 2$ 窗口：
$$\left[\begin{array}{cc}5& 6\\ 8& 9\end{array}\right]$$

$${\text{max}}_{2\times 2}=\max (5,6,8,9)=9$$

3. 生成均值和最大值特征图

均值特征图：

$$\text{mean map}=\left[\begin{array}{cc}3& 4\\ 6& 7\end{array}\right]$$

最大值特征图：

$$\text{max map}=\left[\begin{array}{cc}5& 6\\ 8& 9\end{array}\right]$$

4. 拼接结果

将原始特征图 $x_1$、均值特征图和最大值特征图在通道维度上进行拼接。假设每个特征图表示一个新的通道，则拼接后的特征图 $x_2$ 将有三个通道：

$$x_2=\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]\oplus \left[\begin{array}{cc}3& 4\\ 6& 7\end{array}\right]\oplus \left[\begin{array}{cc}5& 6\\ 8& 9\end{array}\right]$$

如何拼接：填充0
$$\left[\begin{array}{cc}3& 4\\ 6& 7\end{array}\right]\left[\begin{array}{cccc}3& 4& 0& \\ 6& 7& 0& \\ 0& 0& 0& \end{array}\right]$$
$$\left[\begin{array}{cc}5& 6\\ 8& 9\end{array}\right]\left[\begin{array}{cccc}5& 6& 0& \\ 8& 9& 0& \\ 0& 0& 0& \end{array}\right]$$
拼接结果
$$\left[\begin{array}{ccccccccc}1& 2& 3& 3& 4& 0& 5& 6& 0\\ 4& 5& 6& 6& 7& 0& 8& 9& 0\\ 7& 8& 90& 0& 0& 0& 0& 0& 0\end{array}\right]$$

这里的 $\oplus$ 表示在通道维度上拼接。

这种操作使得每个像素都可以“感知”更大范围内的信息，而不仅仅依赖于卷积核的局部感受野。

3.3 改进的跳跃连接和 MLFC 块

问题所在：

在传统的 UNet 中，跳跃连接直接将编码器层的特征传递给解码器层，这可能导致高层次的语义信息无法与低层次的特征充分融合，造成“语义差距”。

MLFC 块的引入：

为了解决这一问题，ACC-UNet 使用了 多级特征编译（Multi Level Feature Compilation, MLFC）块 来改进跳跃连接。MLFC 块将来自不同编码器层次的特征图进行重新调整大小（resize）和拼接，结合多个层次的特征信息，从而减少语义差距。

如何融合特征？

首先，所有编码器层的特征图被调整为相同尺寸，并在通道维度上拼接，然后通过逐点卷积操作减少通道数并融合特征。最终融合后的特征再与对应的解码器层特征图相结合，进一步通过卷积操作增强。

公式表示：

$$x_{\text{comb},i}={\text{PConv}}_{c_{\text{tot}}\to c_i}\left({\text{resize}}_i(x_1)||{\text{resize}}_i(x_2)||{\text{resize}}_i(x_3)||{\text{resize}}_i(x_4)\right)$$
$$x_i={\text{PConv}}_{2c_i\to c_i}\left(x_{\text{comb},i}||x_i\right)$$

举例说明

假设我们有 4 个不同层次的编码器特征图 $x_1,x_2,x_3,x_4$，它们的尺寸分别是：

$x_1:64\times 64\times 32$ （低层次特征图，分辨率最高，通道数最少）

$x_2:32\times 32\times 64$

$x_3:16\times 16\times 128$

$x_4:8\times 8\times 256$ （高层次特征图，分辨率最低，通道数最多）

1. 调整大小和拼接特征图

为了将这些不同大小的特征图结合起来，我们首先需要调整它们的大小使其相同。例如，我们可以将所有特征图调整到 $8\times 8$ 的大小（与 $x_4$ 相同的大小）。

调整大小操作：

${\text{resize}}_4(x_1):x_1\to 8\times 8\times 32$

${\text{resize}}_4(x_2):x_2\to 8\times 8\times 64$

${\text{resize}}_4(x_3):x_3\to 8\times 8\times 128$

$x_4$ 本身已经是 $8\times 8\times 256$，不需要调整大小。

拼接操作：将这些调整大小后的特征图在通道维度上进行拼接：

$\text{concat}({\text{resize}}_4(x_1),{\text{resize}}_4(x_2),{\text{resize}}_4(x_3),x_4)$

拼接后的特征图大小为：

$8\times 8\times (32+64+128+256)=8\times 8\times 480$

2. 逐点卷积总结特征表示

我们使用逐点卷积将拼接后的特征图通道数从 480 减少到与当前层（例如 $x_4$ 的层次）的通道数一致：

逐点卷积操作：

$x_{\text{comb},4}={\text{PConv}}_{480\to 256}(\text{concat})$

结果是一个大小为 $8\times 8\times 256$ 的特征图。

3. 融合和特征增强

将得到的组合特征图 $x_{\text{comb},4}$ 与原始的 $x_4$ 特征图拼接在一起：

$\text{concat}(x_{\text{comb},4},x_4)$

结果大小为：

$8\times 8\times (256+256)=8\times 8\times 512$

逐点卷积减少通道数：使用另一个逐点卷积将拼接后的特征图的通道数从 512 减少回 256：

$x_4={\text{PConv}}_{512\to 256}(\text{concat}(x_{\text{comb},4},x_4))$

最终结果

通过这些步骤，我们得到了一个增强的特征图 $x_4$，它不仅包含了当前层的特征，还结合了其他层次的信息。这种多级信息融合使得特征图具有更强的表示能力，能够更好地捕捉图像的全局和局部信息，从而提升图像分割或其他任务的效果。

4. 性能和评估

ACC-UNet 在五个不同的医学图像分割基准测试中进行了评估，结果显示其性能优于传统卷积网络、Transformer 和混合模型。
在 Dice 系数这一重要指标上，ACC-UNet 相较于 Swin-Unet 和 UCTransNet，分别提高了 2.64±2.54% 和 0.45±1.61%，且参数量仅为它们的 59.26% 和 24.24%。

5. 总结

ACC-UNet 结合了卷积神经网络的归纳偏置和 Transformer 的设计思路，通过 HANC 和 MLFC 块的创新设计，实现了高效的语义分割任务。该模型不仅表现出色，还保持了较低的计算成本，是 2020 年代语义分割领域的一次有力创新。