Restormer: Efficient Transformer for High-Resolution Image Restoration
Syed Waqas Zamir et al., 2022
논문: https://arxiv.org/abs/2111.09881
Github : https://github.com/swz30/Restormer
참고 : https://www.youtube.com/watch?v=3mqu6N4_0pY
※혹시나 포스팅에 문제가 있을 시 댓글 남겨주세요!
[Image Restoration]
Image Restoration 분야는 이미지의 degraded version에서 original image를 복원해내는 것을 목표로 하고 있다. 비가 오는 이미지에서 비를 제거하거나, 흔들림이 있는 이미지에서 흔들림을 제거하는 등의 task를 수행한다.
Image Restoration Task의 approach로는 크게 CNN을 사용한 방법과 Transformer를 사용한 방법으로 구분할 수 있다.
[Image Restoration - CNNs]
- 장점
1) Large-scale dataset로부터 generalizable image prior를 학습하는데 탁월
-Large-scale dataset에 좋은 이유는 아래 그림에서 3*3 크기의 feature map을 생성한다고 했을 때 특정 위치의 weight를 구하기 위해서 convolution(inner product) 연산을 수행하는 데, inner product 연산이 비교적 단순한 연산이기도 하고, transformer처럼 모든 위치의 픽셀을 참조하는 것이 아니라 정해진 개수만큼의 픽셀하고만 계산하기 때문에 large-scale dataset을 사용할 때 빠른 연산이 가능하다고 생각했다. 또한 이미지 데이터의 경우 공간, 구조적인 정보가 중요하다고 생각이 드는데, 아래 그림처럼 인접 픽셀까지 포함하여 weight를 구할 경우 그러한 정보를 잘 잡아내서 generalizable image prior를 학습하는데 수월하다고 이해했다.
- 단점
1) limited receptive field
-위에서 설명한 것처럼 "정해진 개수"의 픽셀과 convolution 연산을 통해 필터의 가중치를 업데이트하기 때문에 large-scale dataset을 사용할 때 빠른 연산이 가능하지만, 그렇기 때문에 receptive field의 크기가 제한되어 먼 거리에 있는 픽셀 간의 연관성을 계산하기 어렵다.
2) static weights at inference
-추론 시에 동일한 convolution filter weight를 사용하기 때문에 입력값에 따라서 유연하게 가중치를 적용하여 결과값을 내는 것이 어렵다.
[Image Restoration - Transformers]
- 장점
1) long-range pixel interaction & High input content adaptability
-Transformer Network의 핵심인 attention module에서는 특정 query마다 task에 더 잘 동작하도록 모든 위치 값을 고려하여 re-weighting해준다. 따라서 먼 거리에 있는 픽셀 간의 연관성을 고려하여 image restoration을 수행할 수 있다. 또한 특정 query마다 더 중요한 가중치를 키우고, 상대적으로 덜 중요한 가중치는 낮추는 방식을 사용하기 때문에 입력값에 따라서 유연하게 결과값을 얻어낼 수 있다. 아래 글에서 attention과 transformer에 관해서 설명하였다.
[논문리뷰] Attention is All You Need
Attention Is All You Need Ashish Vaswani et al., 2017 논문: https://arxiv.org/abs/1706.03762 Github : https://github.com/jadore801120/attention-is-all-you-need-pytorch ※혹시나 포스팅에 문제가 있을..
0choco0.tistory.com
- 단점
1) Complexity
-모든 위치의 픽셀 데이터를 모두 고려해줘야 하기 때문에 H * W image에 대해서 O(W^2 × H^2)의 시간복잡도를 갖게 된다. High resoultion image를 자주 다뤄야하는 image restoration 분야에는 적합하지 않다.
-연산량을 줄이기 위해서 pixel 주변의 small spatial window(8*8)에 self-attention을 적용하거나(https://arxiv.org/abs/2108.10257), 입력 이미지를 48*48 patch로 나누고 각 patch에 독립적으로 self-attention을 적용(https://ieeexplore.ieee.org/document/9577359)한 시도가 있었다. 하지만 self-attetion의 공간적 범위를 제한하는 것은 long-range pixel relationships를 이용하려는 본래 목적에 어긋난다.
[Restormer]
따라서 CNN과 Transformer의 단점을 모두 보완할 수 있는 네트워크가 필요해졌다. 본 논문에서는 적은 연산량으로도 long-range pixel dependencies를 고려하면서 High-resolution image를 처리할 수 있는 Restormer를 제시하였다.
전체적인 네트워크는 위와 같이 전형적인 Encoder-Decoder Transformer 구조를 보이고 있다. 다만 Decoder 중 Transformer Block 하나를 MDTA와 GDFN으로 구현한 Self-attention Block으로 대체하였고, 이것이 Restormer의 핵심이다.
- MDTA(Multi-Dconv Head Transposed Attention)
-기존의 self-attention은 모든 위치의 픽셀을 고려하는 pixel-wise self-attention이기 때문에 앞서 설명했듯이 시간복잡도가 이미지의 height, width의 제곱에 비례하여 커지는 문제가 발생한다. 반면 Restormer의 MDTA는 channel-wise self-attention이다. 즉, pixel 단위에서 attention map을 구하지 않고 channel 단위에서 구하는 것이다. 여기서 핵심 아이디어는 depth-wise convolution이다. 기존의 convolution은 한 개의 convolution layer를 통과하면 한 개의 feature map을 생성한다. 아래 그림을 보면, 입력 데이터의 사이즈는 (Channel, Height, Width) = (3, 8, 8), convolution layer의 사이즈는 (3, 3, 3)이고 둘의 convolution 연산을 통해 (3, 8, 8) feature map 중 한 개의 weight를 결정하고 있다.
Depth-wise convolution은 (3, 3, 3) 크기의 convolution layer를 채널별로 나누어 convolution 연산을 수행한다. 즉, (1, 3, 3) 크기의 convolution layer 3개가 아래 그림처럼 생성되고 각기 다른 feature map을 생성하게 된다.
Restormer의 MDTA는 아래와 같다. 먼저 마지막에 더해줄 residual X를 빼놓고, X를 layer-normalization해준다. 이후 depth-wise convolution을 사용하여 채널별로 생성된 feature map을 통해 channel-domain에서 feature covariance를 계산하서 attention map을 얻어낸다. 기존의 pixel-wise가 아닌 channel-wise self-attention을 수행하여 계산 복잡도를 줄였다. High-resolution image의 경우 H, W보다는 C가 작기 때문이다. 따라서 적은 연산량으로도 long-range pixel dependencies를 활용한 self-attention을 구현한 것이다. 또한 depth-wise convolution도 convolution 연산의 일종이기 때문에 기존 CNN의 장점인 local한 공간 정보또한 얻어낼 수 있다.
- GDFN(Gated-Dconv Feed-Forward Network)
GDFN은 기존 attention module의 feed-forward와 유사하지만 gating mechanism을 추가하여 두 개의 1*1 conv layer와 3*3 depth-wise conv layer를 통과시키게 된다. 이때 γ배만큼 채널 수를 늘리고 아래쪽에는 Gelu(Gaussian error linear units) function을 적용하여 모델의 복잡도를 증가시켰다. 이후 1*1 conv layer를 통과시켜 이미지의 채널 수를 맞춰주고 residual(X)을 더한다. 이 과정을 통해 encoder-decoder pipeline의 각 level(H, W는 감소하고 C는 증가하면서 네트워크가 깊어지는 것)이 다른 level과 보완하며 task에 적합한 정보에 집중하게 된다.
[Experimental Results]
아래와 같이 Image Restoration의 de-rain, de-blur, de-noise task에 대해서 Restormer가 SOTA 성능을 보이고 있음을 확인할 수 있다.
아래는 연산량에 대한 ablation study인데, MDTA+GDFN을 사용했을 때 연산량이 실제로 감소하지는 않아서 확인해볼 필요가 있어 보인다. 하지만 연산량의 차이가 그렇게 크지 않은데, 기존의 Resblock으로 구성된 UNet baseline model보다 성능이 더 좋다는 점에 의의를 두는 것 같다.
'딥러닝' 카테고리의 다른 글
| [논문리뷰] Attention is All You Need (0) | 2022.07.31 |
|---|---|
| [논문리뷰] Specular Reflection Separation using Dark Channel Prior (0) | 2022.07.28 |
| [논문리뷰] Accurate Image Super-Resolution Using Very Deep Convolutional Networks (0) | 2022.07.24 |
| [논문리뷰] Image Super-Resolution Using Deep Convolutional Networks (0) | 2022.07.24 |
