arXiv:2412.18276cs.CVeess.IV2024-12中稿 · ACCV2024被引 1

提出轻量版U-Net,内存减少93.3%且性能提升

UNet--: Memory-Efficient and Feature-Enhanced Network Architecture based on U-Net with Reduced Skip-Connections

论文配图:UNet--: Memory-Efficient and Feature-Enhanced Network Architecture based on U-Net with Reduced Skip-Connections
图 1 · 摘自论文原文
  • 用多尺度聚合模块减少编码器内存占用
  • 解码器增强模块恢复并提升特征表示能力
  • 适配图像修复等任务,通用性强

基于编码器-解码器结构与跳跃连接的U-Net模型在多种视觉任务中表现优异。跳跃连接需在解码前保持编码器特征图,导致内存开销大,不利于资源受限设备。本文提出一种通用方法与架构,通过在编码器设计多尺度信息聚合模块(MSIAM),将多尺度特征压缩为单尺度以降低内存占用;再在解码器引入信息增强模块(IEM),将压缩特征扩展并重构为多尺度特征,实现高效且增强的特征表达。将该方法应用于图像修复领域SOTA模型NAFNet,构建出内存高效的UNet--架构。实验显示,其跳跃连接内存需求降低93.3%,同时性能优于NAFNet。此外,该方法可泛化至多个视觉任务,在降低内存消耗的同时,提升网络精度,优于现有高效架构。

原文摘要 · Abstract (English)

U-Net models with encoder, decoder, and skip-connections components have demonstrated effectiveness in a variety of vision tasks. The skip-connections transmit fine-grained information from the encoder to the decoder. It is necessary to maintain the feature maps used by the skip-connections in memory before the decoding stage. Therefore, they are not friendly to devices with limited resource. In this paper, we propose a universal method and architecture to reduce the memory consumption and meanwhile generate enhanced feature maps to improve network performance. To this end, we design a simple but effective Multi-Scale Information Aggregation Module (MSIAM) in the encoder and an Information Enhancement Module (IEM) in the decoder. The MSIAM aggregates multi-scale feature maps into single-scale with less memory. After that, the aggregated feature maps can be expanded and enhanced to multi-scale feature maps by the IEM. By applying the proposed method on NAFNet, a SOTA model in the field of image restoration, we design a memory-efficient and feature-enhanced network architecture, UNet--. The memory demand by the skip-connections in the UNet-- is reduced by 93.3%, while the performance is improved compared to NAFNet. Furthermore, we show that our proposed method can be generalized to multiple visual tasks, with consistent improvements in both memory consumption and network accuracy compared to the existing efficient architectures.

U-Net内存优化图像修复特征增强

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。