目录
- 1. Introduction
- 2. Method
- 2.1. Overview
- 2.2. UFONE
- 2.3 真实场景下的部署优化
- 3. 结果
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的主要思路是：
- unfold和fold操作只在网络的开始和结束阶段进行一次，因为去除了cyclic-shift, 提出Inner-patch Transformer Layer（ITL）， self-attention只在window内部进行；
- 提出Spatial-Aware Layer (SAL)，通过空间注意力，来感知长距离依赖，来弥补ITL忽略掉的长距离信息；
- 在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,
