ACM-MM2023 DITN详解:一个部署友好的超分Transformer

Paper: Unfolding Once is Enough: A Deployment-Friendly Transformer
Unit for Super-Resolution

Code: https://github.com/yongliuy/DITN

1. Introduction

• CNN做超分的缺点
  由于卷积权重在推理时是固定的, 而且在整个图像上共享，无法动态提取不同输入以及不同局部区域的特征；

• Vision transformer做超分的演化

  SwinIR将swin transformer应用到超分任务中 --> ELAN提出了一些效率优化改进方案，并提出了一个更有效的感知长距离依赖的方法

  但是由于cyclic-shift的存在，这些方法仍然需要对feature map反复进行fold和unfold操作，尽管这些reshape算子不会被计算在FLOPs和parameter中, 但是在推理阶段会计算负担和显存占用，尤其是在做dynamic inputs的部署时。

本文提出Deployment-friendly Inner-patch Transformer Network (DITN) 减少reshape算子的使用，DITN的主要思路是：

1. unfold和fold操作只在网络的开始和结束阶段进行一次，因为去除了cyclic-shift, 提出Inner-patch Transformer Layer（ITL）， self-attention只在window内部进行；
2. 提出Spatial-Aware Layer (SAL)，通过空间注意力，来感知长距离依赖，来弥补ITL忽略掉的长距离信息；
3. 在TensorRT部署加速阶段，替换掉了LN层，提出了一种融合优化策略，进一步减少了显存占用和计算负担；

2. Method

2.1. Overview

作者的设计动机是图像分类和目标检测这些high level的任务，通常需要从全局特征中获取大部分的语义信息，而超分这种low-level的任务同样需要去探索局部特征去生成精细的细节和纹理。

算法的整体结构如下图所示，延续了SwinIR的“浅层特征提取+深层特征提取+图像重建”的结构，区别就是在深层特征提取模块，使用了新提出的A Deployment-Friendly Transformer Unit（UFONE）替代RSTB。

2.2. UFONE

UFONE由两个基本的层组成，一个是Inner-patch Transformer Layer (ITL)，另一个是Spatial-Aware Layer (SAL)。ITL和SAL是串联的；

ITL完成的是window内部的特征提取，其实就是w-msa，SAL完成的是长距离信息的感知，因为ITL忽略了这部分信息。下面我们来看一下他们分别是怎么实现的。

• Inner-patch Transformer Layer (ITL)
  
  其实就是一个w-msa，唯一的不同就是完成了self-attention之后将mlp换成了一个卷积

• Spatial-Aware Layer (SAL)

首先输入整幅图X，在经过一个卷积层后在通道维度分成X1和X2，随后X1在经过空洞卷积后与X2相乘，完成了空间上的注意力，从而实现了长距离信息的感知。

2.3 真实场景下的部署优化

• Layer norm用tanh和1x1卷积替代，layer norm和tanh都能把输入的值映射到[-1, 1]之间，1x1卷积可以完成线性变换
  这一改动会轻微降低模型效果，提速8%~10%；

• self-attention层的算子优化

  （1）算子融合；
  

  （2）FlashAttention

  cuda实现后通过tensorrt的plugin机制使用；

3. 结果

DITN: 3个UFONE
DITN-Tiny：1个UFONE
DITN-Real：DITN-Tiny的基础上替换掉layer norm
可以看出比ELAN稍快。此外swinir和elan不支持动态尺寸输入，所以在动态输入上没有进行比较，文中是这么说的：

Noting that the unique operators in ELAN-Tiny and SwinIR-Tiny make them
unavailable in the dynamic input configurations, we compare only DITN-Tiny and DITN-Real. Specifically,