SEMANTIC SEGMENTATION IN LEARNED COMPRESSED DOMAIN (arxiv 2022)

2022/09/14:今天是八月深圳疫情后正常上班第一天，回到了每天+350和免费健身房的日子，能够感到快乐。

paper：https://arxiv.org/abs/2209.01355

Abstract

大多数机器视觉任务(如语义分割)都是基于图像压缩算法(如JPEG)编码和解码的图像。然而，这些在像素域内解码的图像引入了失真，并且它们是针对人类感知进行优化的，这使得机器视觉任务的性能处于次优状态。在本文中，我们提出了一种基于压缩域的方法来改进分割任务。i)提出了一种动态和静态信道选择方法，以减少编码获得的压缩表示的冗余。ii)探索和分析两种不同的变换模块，帮助将压缩的表示转化为分割网络中的特征。实验结果表明，与最先进的基于域的压缩工作相比，我们可以节省高达15.8%的比特率，与基于像素域的方法相比，节约比特率为83.6%，推理时间为44.8%。

1. Introduction

随着深度学习的发展，机器视觉任务在过去的几年里得到了广泛的实现，并在我们的生活中扮演了重要的角色。例如，语义分割被应用于医学图像分析、自动驾驶汽车、视频监控、增强现实等许多领域。然而，许多机器视觉任务通常是基于解码图像，需要图像压缩。解码后的图像失真影响了机器视觉任务的执行性能。

经典的图像压缩算法，如JPEG, HEVC / H.265-intra在人类感知中取得了良好的率失真性能，而一些基于学习的端到端图像压缩方法也取得了可竞争的性能。然而，所有这些设计都是为了减少人类感知到的失真，与机器感知有差距。因此，目前的算法通常会导致次优的语义分割。

针对这些问题，提出了一些基于压缩域的方法。基于像素域的方法意味着我们需要将经过图像压缩算法编码解码的图像输入到语义分割网络中，如图1所示。与此不同的是，基于压缩域的方法是直接将压缩域的表示形式输入到语义分割网络中，而不需要对解码后的图像进行计算。提出了相关标准JPEG AI。Torfason等已经证明了基于压缩域的方法比基于像素域的方法能达到更好的性能，并在图像分类和目标检测上做了一些实验。Choi等人通过使用信息论方法进一步验证了这一点，并提出了基于压缩域的方法，该方法在对象检测和语义分割上得到了验证。Bai等利用Transformer改进了基于压缩域的图像分类。Wang等提出了信道选择方法和熵估计方法来改善[7]的结果。然而，信道选择方法只能删除熵最小的信道，因为信道选择方法是基于压缩表示的方差。

在本文中，我们提出了一种基于压缩域的方法来改进分割任务。本文的主要贡献如下:

• 在学习门模块的基础上，提出了一种动态和静态信道选择方法，以选择合适的信道并减少压缩表示的冗余。
• 我们探索和分析了两个不同的模块，将压缩域中的表示转换为来自分割网络中间层的特征。
• 我们进行了基于cityscape的实验，结果表明，在相同的mIoU情况下，与最先进的基于域的压缩方法相比，我们可以节省15.8%的比特率，而与基于像素域的方法相比，我们可以节省约83.6%的比特率和44.8%的推断时间。

2. Method

2.1 Framework

本文方法的框架如图2所示。其中包括五个模块：一个图形压缩网络，一个门控模块，一个变换模块，一个分割网络和一个知识蒸馏链接。文献[4]中的超先验模型被设置为本文的图像压缩网络，其预训练参数由Compress AI提供。在后续步骤中，压缩网络的权重被固定，因此我们的方法不会影响解码图像的质量。

对于语义分割网络，我们选择DeepLab v3[12]，去掉前两层以符合压缩表示的大小。然后，为了获得更好的性能，采用了知识蒸馏的方法。知识蒸馏的教师网络显示在包含DeepLab v3前两层的知识蒸馏路径中，通过原始图像进行预训练，之后进行冻结。学生网络是知识蒸馏链接上方的所有网络结构。
我们首先将原始图像x输入到压缩网络的编码器中，得到压缩后的表示形式$\hat{y}$。对于基于像素域的方法，我们需要将这种表示形式$\hat{y}$传输到解码器中，得到解码后的图像$\hat{x}$，并将$\hat{x}$输入到语义分割网络中完成分割。但是对于我们的基于压缩域的方法，为了得到机器视觉任务的结果，我们不需要解码器来获取$\hat{x}$，只需要向门控模块输入$\hat{y}$就可以得到一个选中的表示${\hat{y}}_{selected}$。

然后利用变换模块将$y_{selected}$变换为分割网络的特征$\mathcal{F}$。再将特征输入到语义分割网络中，得到分割映射$\mathcal{M}$。计算$\mathcal{M}$和地面真值图像${\mathcal{M}}_{gt}$的交叉熵函数${\mathcal{L}}_{seg}$用于更新网络。可以将其表示为：
其中$\mathbb{SG}$为停止梯度算子，因为在训练过程中压缩网络参数$\phi$没有更新。为了进一步提高性能，利用知识蒸馏链接计算另一个知识蒸馏损失${\mathcal{L}}_{KD}$。我们将$x$输入到冻结的语义分割网络的前两层，得到${\mathcal{F}}_{gt}$。然后计算${\mathcal{F}}_{gt}$和$\mathcal{F}$的均方误差为${\mathcal{L}}_{KD}$，可以表示为：

其中$\psi$表示知识蒸馏链接的权值，这些权值在训练过程中被冻结。最后，总损失L可表示为:

其中${\lambda }_1$用于控制${\mathcal{L}}_{KD}$和${\mathcal{L}}_{seg}$之间的权重，这里根据辅助实验将${\lambda }_1$设为1。另一方面，通过继续传输$\hat{y}$和${\hat{y}}_{selected}$之间的残差，我们也可以得到解码后的图像。

2.2 Channel Selection

先前的工作[8,10,13]表明压缩表示中有一些通道对于机器视觉任务是无用的。在此，我们提出了一个门控模块来自适应学习哪个通道适合于语义分割任务。在此基础上，提出了动态选择和静态选择方法。

2.2.1 Gate Module

门控模块设计如图3所示。当我们得到压缩表示$\hat{y}$，我们将其输入到门控模块中得到挑选后的表示${\hat{y}}_{selected}$。因为在信道选择的过程中，我们需要对信道进行采样，这可能会使网络不可微。因此，与[14]类似，引入gumbel-softmax层，使整个门模块可微。

2.2.2 Dynamic Selection

直接在门控模块中输入表示$\hat{y}$可以使我们动态地选择通道。例如，我们假设$\hat{y}$的大小是$H\times W\times C$。经过池化层和两个卷积层后，我们得到了一个$1\times 1\times 2C$的张量，其中第一个$C$数表示压缩表示$\hat{y}$的$C$个通道分别被采样为$0$的概率，而后一个$C$数表示$C$个通道被采样为$1$的概率。对于最后的Gumbel-Softmax层，我们可以得到一个$1\times 1\times C$的张量，其中每个元素的值为0或1，表示我们没有选择或选择相应的通道。最后，我们将这个张量乘以$\hat{y}$，得到${\hat{y}}_{selected}$。整体流程如图3所示

2.2.3 Static Selection

然而，动态选择会使每个图像选择不同的通道，这可能会导致训练的不稳定性。因此，在动态选择的基础上，我们进一步提出了一种静态选择方法。根据动态选择的方法训练整个网络后，我们使用训练数据集中的所有图像进行推断，并计算每个通道被动态选择的次数。然后，我们选择统计中最频繁被选择的$N$个通道。最后，我们手动选择这些通道，而不是使用门控模块来获得${\hat{y}}_{selected}$。

2.3 Transform Module

得到$\hat{y}$后，我们将其输入到一个变换模块中，因为压缩表示的大小与分割网络中特征的大小不同，压缩域与分割网络中的域存在差距。为了有效地弥合这些差距，并将压缩表征${\hat{y}}_{selected}$转换为特征$\mathcal{F}$用于分割。我们测试了两种不同的变换模块。

针对[4]，我们首先设计了包含两个转置卷积层的Module-Deconv，如图4(a)所示，[13]中也使用了该模块。为了提高性能，根据[5]，Module-Res的设计如图4(b)所示。残差块和PixelShuffle用于提高性能。

3. Experiments

3.1 Experiment Setting

对于图像压缩网络，我们使用文献[4]中的超先验模型。对于语义分割网络，我们使用
DeepLab v3。Cityscenes数据集用作训练集和测试集。Adam被选为优化器。我们总共训练这个网络400 epoch，初始学习率设置为0.001,300 epoch后下降到1e-4。所有实验都基于GeForce RTX 3090 24 GB GPU。

3.2 Rate-mIoU Comparison

我们测试了动态和静态的选择方法,在本节中基于Module-deconv。在静态选择方法的实验中，我们分别测试了选择通道数为128、64、32和16的情况。

为了证明压缩表示中确实存在冗余，我们将我们的方法与不使用任何通道选择方法的情况进行了比较。如图5所示，与这种情况相比，我们最多可以节省31.8%的比特率。

为了验证该方法的有效性，我们首先将其与基于像素域的压缩网络同样为[4]的方法进行了比较。如表1所示，我们基于压缩域的方法优于基于像素域的方法。此外，我们将我们的方法与最先进的基于域的压缩方法[10]进行比较。在这里，我们使用我们的方法中相同的图像压缩网络来再现[10]。如图5所示，无论选择多少通道，我们的静态方法总是优于[10]。例如，如果选择的通道数量是32，相比[10]我们可以节省高达15.8%的比特率，与基于像素域的方法相比节省约83.6%。这表明了我们的方法可以选择更适合分割任务的通道。当所选通道小于128时，静态方法具有较好的rate-mIoU，表明静态方法优于动态方法。

3.3 The Rate-PSNR Comparsion

由于在训练过程中图像压缩网络的权值被冻结，因此其rate-psnr与文献[4]的结果相同。同时，我们再现了信道选择方法[10]通过使用与我们相同的图像压缩网络，这意味着我们有相同的rate-psnr性能。

3.4 The Effect of Various Channels

在图6中，我们测试了各种通道的用例。我们可以看到，与选择64或128个通道相比，选择
16或32个通道可以得到更好的rate-miou，这意味着我们减少了压缩表示中的冗余。

3.5 The Effect of Various Transform Modules

在本文中，我们探讨了两个不同的变换模块。如表2所示，Module-Res在两种比特率下都能得到更好的rate-mIoU，而Module-Deconv的表现稍差，但其复杂度只有Module-Res的一半。

3.6 Inference Time

为了检查我们方法的效率，我们基于Cityscenes数据集测试推理时间，如表3所示。“No Gate”是指不使用任何信道选择的基于压缩域的方法。可以看到，与基于像素域的方法相比，压缩的基于域的方法可以节省大量的时间。此外，使用动态选择和静态选择时，我们可以节省更多时间，特别是选择通道数为16时，我们可以节省大约44.8%的时间。

3.7 Scalability for Multi-tasks

由于图像压缩网络的权值是固定的，因此该方案具有多任务可扩展性。通过使用基于PASCAL VOC 07数据集的动态选择方法测试对象检测，我们探索了潜在的可伸缩性。使用Fasterrcnn[16]作为任务网络。结果表明，与基于像素域的方法相比，我们节省了约30%的比特率，达到了相似的平均精度。进一步探索可伸缩性是我们未来的工作。

4. Conclusion

在本文中，我们提出了一种基于压缩域的方法来改进分割任务。我们首先提出了一个门控模块来自适应学习压缩表示中的哪些通道适合分割任务。在此基础上，研究了指定通道数量的静态选择和基于门模块的动态选择。这两种选择帮助我们成功地减少了压缩表示的冗余。然后，引入两种不同的变换模块，将压缩域中的表示转换为分割网络中的特征。最后，我们在城市景观数据集上进行了实验，实验结果表明，我们的模型比之前的方法能够获得更好的rate-miou和更快的推理速度。