通过分组移位窗口注意力,实现图像修复模型的低内存高效率训练。
AgileIR: Memory-Efficient Group Shifted Windows Attention for Agile Image Restoration
- 将自注意力分解为多组,减少反向传播内存占用。
- 在Set5上保持32.20 dB性能,大批次下节省超50%内存。
- 适合追求高效训练与低显存需求的图像修复研究者。
图像Transformer在图像修复任务中表现卓越,但多数基于Transformer的模型因内存占用过高而受限。本文旨在降低Swin Transformer的内存消耗,并加快训练速度。为此提出AgileIR,结合分组移位窗口注意力机制与窗口注意力,从结构上稀疏化模型。我们设计了分组移位窗口自注意力(GSWA),将移位窗口多头自注意力(SW-MSA)和窗口多头自注意力(W-MSA)按注意力头分组,有效降低反向传播时的内存开销。同时,在训练中保留移位窗口掩码及其可学习偏置,以促进跨窗口通道交互。此外,重新分配投影参数以加速注意力矩阵计算,仅导致可忽略的性能下降。实验表明,相较于基线SwinIR及其他高效量化模型,AgileIR在大批次设置下仍保持32.20 dB的峰值信噪比(PSNR),优于其他定制化高效方法,内存消耗减少超过50%。
原文摘要 · Abstract (English)
Image Transformers show a magnificent success in Image Restoration tasks. Nevertheless, most of transformer-based models are strictly bounded by exorbitant memory occupancy. Our goal is to reduce the memory consumption of Swin Transformer and at the same time speed up the model during training process. Thus, we introduce AgileIR, group shifted attention mechanism along with window attention, which sparsely simplifies the model in architecture. We propose Group Shifted Window Attention (GSWA) to decompose Shift Window Multi-head Self Attention (SW-MSA) and Window Multi-head Self Attention (W-MSA) into groups across their attention heads, contributing to shrinking memory usage in back propagation. In addition to that, we keep shifted window masking and its shifted learnable biases during training, in order to induce the model interacting across windows within the channel. We also re-allocate projection parameters to accelerate attention matrix calculation, which we found a negligible decrease in performance. As a result of experiment, compared with our baseline SwinIR and other efficient quantization models, AgileIR keeps the performance still at 32.20 dB on Set5 evaluation dataset, exceeding other methods with tailor-made efficient methods and saves over 50% memory while a large batch size is employed.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。