뭉그적뭉그적

오늘 소개할 논문은 cyclegan을 활용해 segmentation task에서 domain adaptation을 진행한 논문이다. 

먼저 cyclegan을 이용해 source domain을 target domain으로 domain adaptation을 진행하고 target domain으로 스타일이 바뀐 source data로 segmentation 네트워크를 학습하여 실제 target domain에서의 성능을 높이는 방법이다. 

즉 전체 네트워크는 2 step이다. 

1. source data의 domain을 target data의 domain으로 바꾸기 (cyclegan)

2. target domain으로 옮겨진 source data'을 이용하여 segmentation 네트워크 학습시키기 (deeplab)

이런 2 step domain adaptation은 사실 이전에도 여러번 소개가 되었다. 

Dual Channel-wise Alignment Networks for Unsupervised Scene Adaptation(DCAN) 이나 cyCADA 논문도 위와 같은 형식이라고 알고 있다. 

그런데 이렇게 2 step으로 네트워크를 학습시킬 때 발생하는 문제점이 있다. 첫번째 step에서 Image Generator의 성능이 떨어진다면, 두번째 step에서의 성능이 결코 좋을 수 없다는 것이다. 즉, segmentation network의 성능이 image generator의 성능에 의해 한계점을 지니게 된다. 

본 논문에서는 위와 같은 한계점을 보완하는 방법을 제시한다. 

Contribution

논문에 제시되어 있는 자신의 논문의 기여는 다음과 같다고 한다.

"""

1. Bidirectional Learning System, which is a closed-loop to learn the segmentation adaptation model and the image translation model and the image translation model alternatively

2. Self-supervised learning algorithm for the segmentation adaptation model, which incrementally align the source domain and the target domain at the feature level, based on the translated results

3. New Perceptual loss to the image-to-image translation, which supervised the translation by the segmentation adaptation model

"""

1. 2step을 개별적인 단계가 아닌 유기적으로 이어지게 (closed-loop) 만듦으로써 앞뒤 스텝의 단점을 상호 보완 가능하게끔 만들었으며,

2. 두번째 스텝에서 self-supervised learning algorithm을 추가하였고,

3. 첫번째 스탭에서 두번째 스텝 네트워크를 활용한 perceptual loss를 추가하여 segmentation adaptation model로부터 training에 필요한 정보를 얻었다고 한다. 

Model 

전체 네트워크는 위와 같다. 

Image Translation Model을 통해 source data를 target domain으로 shift 한 후, 그 shift된 source data를 Segmentation Adaptation Model을 학습하여 결과를 얻는다. 그런데 Image Translation Model에서 Segmentation Adaptation Model 아웃풋을 이용한 loss인 perceptual loss가 추가된다. Segmentation Adaptation Model에서는 adversarial loss가 존재하고, 스스로 mask를 만들어 self-supervised learning을 진행한다고 하는데... 자세한 내용은 loss를 천천히 뜯어보면 알 수 있겠다. 

Loss

loss들을 같이 보기 전, 기억해야할 notation은 다음과 같다. 

source data = S

target data = T

translated source data by forward cyclegan = S'

translated target data by inverse cyclegan = T'

Image Translation Model (cyclegan generator) = F

cyclegan inverse generator = F-1

Segmentation Adaptation Model = M

첫번째 스텝인 image translation model로 cyclegan을 활용하였으며,

그 loss는 

이다. 람다 GAN은 1, 람다 recon은 10 을 사용했다고 한다. 

하나하나 뜯어보자. 첫번째로 GAN loss이다. 

S'과 T의 domain을 구분하지 못하도록. 또는 S와 T'의 domain을 구분하지 못하도록 만드는 loss이다. discriminator는 S'을 0으로, T 를 1로 구분하도록 학습하였으니, 반대로 generator는 S'을 1로, T를 0으로 라벨하여 loss를 얻는 것을 볼 수 있다. 

F를 통과하여 만들어진 S'을 다시 F-1 함수에 통과하였을 때, S로 되돌아오는지, reconstruction 되는지 확인하는 pixel-level loss이다. L1 loss를 사용하고 있다. cyclegan 논무에 따르면, 이 loss를 추가함으로써 generator F가 다시 되돌릴 수 있을 만큼만 이미지를 바꿔주어 이미지의 content를 보존할 수 있었다고 한다. 

논문의 핵심 특징 중 하나인 perceptual loss이다. 두번째 step의 segmentation model과 연결시켜주는 고리이기도 하다. segmentation adaptation model M을 마치 perceptual loss의 VGG 네트워크처럼 활용하여 S와 S'의 perceptuality, 그리고 F-1(S')과 S의 perceptuality를 줄이는 방향으로 학습한다. 람다 per은 0.1, 람다 per_recon은 10이다. 이러한 perceptual loss를 추가함으로써 segmentation adaptation model M이 가지고 있는 domain bias를 image translation model에서 보완해 줄 수 있게 된다. 

두번째 스텝인 segmentation adaptation model로 deeplab v2를 사용하였으며,

그 loss는 self-sueprvised loss의 유무로 두가지 종류가 존재한다. 

람다adv는 0.001이다. 

이번에도 하나하나 뜯어보자. 

먼저 위 논문의 큰 특징 중 하나인 adversarial loss이다. image translation model 뿐만이 아닌 segmentation model에도 adversarial loss를 추가함으로써 image translation model에서 미처 줄이지 못한 source와 target 사이의 거리를 줄이는 방향으로 보완해 준다. 

segmenatation model에서 항상 등장하는 cross entropy loss이다. domain shift 된 S'을 input으로 하여 얻은 M(S') 결과값과 segmentation label Ys 사이의 loss를 계산한다. 

마지막으로 핵심 loss인 self-supervised segmentation loss이다. target data label이 주어지지 않았을 때, 스스로 라벨을 형성하여 unsupervised learning을 진행한다.

학습이 덜 된 네트워크에서 나온 결과값을 정답값이라고 가정한다는 것이 아이러닉하게 들릴 수도 있겠다. 그러나, segmentation task의 특징을 잘 생각해 보자. 도로를 segmentation한다고 했을 때, 도로의 정 중앙이 도로에 속한다는 것은 쉽게 알 수 있지만, 도로의 가장자리, 즉 나무 또는 보도와 가까운 pixel일수록 도로인지 아닌지 판단하는 것은 어려워진다. 즉, 같은 segmentation task에서도 pixel마다 쉽게 class를 판단 가능한지 아닌지 난이도의 차이가 존재하기에 확률의 차이가 존재하게 된다.

따라서 높은 확률로 판단 가능한 쉬운 pixel에 대해서 우리는 sudo label을 만들 수 있다. 몇퍼센트의 확률로 class를 확신하면 sudo label로 설정할 것이지에 대한 threshold는 0.9로 정했다고 한다. 이렇게 정해진 threshold로 우리는 target data에 대한 sudo label을 구할 수 있고, 이 sudo label을 논문에서는 mask라고 부른다. figure 3에서 예시로 주어진 mask를 살펴보면 구분선은 검정색(0)으로 구분하기 어려운 곳이기 때문에 sudo label을 형성하지 못한것을 볼 수 있고, 특정 class의 내부는 흰색(1)로 구분하기 쉬웠기 때문에 sudo label이 형성되었음을 볼 수 있다. 이렇게 정해진 sudo label에 대해서만 target segmentation loss를 구하여, training을 더욱 가속화한다.

Training Process

이렇게 정의된 2 step network를 구성하는 F,M model은 다음과 같은 프로세스로 training이 진행된다. 

먼저 일반 segmentation model과도 같은 M(0)를 학습시킨다. 위의 M을 training 하는 loss에서 adversarial loss와 self-supervised segmentation loss를 제외한 segmentation loss만을 활용한 것을 말한다.

이 M(0)를 기반으로 cyclegan F(1)을 학습한다. <- Equation 2

F(1)의 아웃풋으로 adversarial loss와 segmentation loss를 갖는 M을 학습한다. <- Equation 1

그 후 M에 self-supervised segmentation loss를 추가하여 M을 N번 학습한다. <- Equation 3

위와 같은 과정을 K번 진행했다고 한다. K는 2, N은 3으로 최종 네트워크를 지정하였다. 

이렇게 논문의 전반적인 내용이 끝났다. 결과는 생략하겠다. 

이 논문에서 인상 깊었던 것은 

1. 2 step 네트워크일 때 두 네트워크를 유기적으로 연결한 것

2. segmentation model에서 sudo label을 만들어 self-supervised learning을 활용한 것

이다. 두 특징은 segmentation task에만 한정이 되는 것이 아니라 같은 idea로도 다양한 task에 활용될 수 있을 것 같다. 

이번 포스팅에서는 2019년 ICCV에서 발표된 wavelet style transfer 이용한 super-resolution 논문을 리뷰해 보려 한다.

0 Abstract 부터 2 Related work는 배경설명 및 다른 네트워크들간의 비교이며 본 논문의 모델은 3 proposed method에서 소개한다. 

<간단 요약>

1. EDSR, CX를 이용하여 성질이 다른 두 이미지를 얻는다. 

2. 두 이미지를 stationary wavelet transform(SWT)을 통해 low frequency sub domain과 high frequency sub domain을 구분한다. 

3. high frequency sub domain를 pixel domain이 아닌 wavelet domain에서 style transfer를 함으로써 high frequency perceptuality를 개선하였다. 

4. EDSR의 low frequency sub domain을 압축된 VDSR에 통과시켜 objective quality를 높였다. 

5. content loss와 style loss 모두를 고려하는 새로운 loss function을 제안했다. 

6. 그 결과 high frequency detail이 살아났으며 선명한 텍스쳐와 구조를 얻을 수 있었다. 즉, PSNR과 NRQM(perceptuality)를 모두 높일 수 있었다. 

0 Abstract

super-resolution은 LR(low-resolution, 저해상도) 이미지로부터 HR(high-resolution, 고해상도) 이미지를 복원하는 것을 말한다. 최근 super-resolution 부문의 이슈는 low distortion과 high perceptual quality가 trade-off 관계라 두 관점을 모두 만족시키기 어렵다는 것이다. 

지난번에 언급한 SRCNN, VDSR, EDSR 등의 논문들은 모두 LR이 HR에 최대한 가까워 지도록 학습해 distortion은 적지만(PSNR은 높지만) 엣지가 선명하지 않고 디테일이 사라지며 질감이 무너지는 것과 같이 사람의 눈으로 보았을 때 만족스럽지 않은(low perceptual quality) 결과를 얻었다. 

이런 low perceptual quality 문제점을 최근 GAN 을 기반으로 한 모델이 photo realistic한 이미지를 만들어 내며 잘 보완해 주고 있다. 그러나 이 논문에서는 wavelet domain style transfer(WDST)를 이용하여 GAN을 기반으로 한 모델보다 perception-distortion(PD) trade-off 를 더 잘 보완할 수 있었다고 한다.

2D stationary wavelet transform(SWT)를 이용하여 이미지의 low-frequency 부분과 high frequency 부분을 나눠 각각에 대해 독립적으로 개선시킨 후 합친다는 것이 이 논문의 취지이다. low-frequency sub-bands를 이용하여 objective quality를 개선하고(low distortion), high-frequency sub-bands를 이용하여 perceptual quality를 개선할 수 있다. 

1 Introduction  &  2 Related work

앞서 말했듯이 기존의 연구는 mean square error(MSE)를 낮춰 objective image quality를 높이는 데 집중한 연구와 adversarial training을 통해 perceptual loss를 낮춰 perceptual image quality를 높이는 데 집중한 연구로 나뉜다. 아래의 그림1을 보면 PSNR을 높이는 데 집중한 파란색 네트워크들과 perceptuality(NRQM metric)를 높인 초록색 네트워크들이 명확히 구분된다.

objective quality를 높인 그림1의 파란색 네트워크들은 대부분 MSE를 minimize하여 reconstructed image와 ground truth image 사이의 간극을 좁힌다. 이는 high frequency detail들이 사라져 엣지가 blur된다는 문제점을 야기한다. 

perceptual quality를 높인 그림1의 초록색 네트워크들은 VGG loss와 adversarial loss의 weighted sum으로 구성된 perceptual loss를 낮추도록 학습한다. VGG loss는 reconstructed image와 ground truth image 간의 perceptual similarity를 높이는 데 효과적이며, adversarial loss는 reconstructed image를 realistic하게 만든다.

위 trade-off 관계를 보완하기 위하여 adversarial loss와 MSE를 활용하는 SRGAN-MSE, ENet 등의 그림1의 주황색 모델이 도입되었으나 unstable하다는 지적이 있었다.  SRGAN-MSE는 MSE loss와 adversarial loss를 합쳐 사용하는데, adversarial loss는 결과 이미지의 high frequency detail들을 살려냈지만, 그것이 정확한 위치에 있지 않아서(ground truth image와는 사뭇 다른 이미지여서) MSE distortion은 증가되는 문제점이 있었다. 이런 instability 문제점을 보완하기 위해 texture matching loss를 도입하였지만 여전히 blocking과 noisy artifact가 생겼다. 

원치 않는 artifact들이 생기는 문제점을 보완하기 위하여 ESRGAN이 등장하였다. PIRM challenge에서 우승한 ESRGAN은 MSE를 낮추는 프로세스와 perceptual quality를 높이는 프로세스를 분리하기 시작한다. 두개의 독립적인 네트워크는 마지막에 interpolation하여 합쳐진다. 그러나 network interpolation은 두 네트워크가 완전히 똑같은 구조를 가지고 있어야 하기에, 분명히 다른 성질을 갖고있는 두 부분에 대해 각각의 성장 가능성을 저해한다는 단점이 있다. 

따라서 network interpolation이 아닌 분리된 네트워크에서 나온 image들의 fusion을 고안해 보았을 때, 네트워크 구조의 유연함을 증폭할 수 있다는 점이 큰 장점이다. 최근 Deng 이 style transfer를 이용한 두 이미지의 fusion을 제시하였다. 그러나 style transfer는 pixel domain상에서의 작업이므로 이미지의 structure와 texture를 동시에 보존한다는 것은 매우 까다롭다. 그림2를 보면 objective and perceptual quality 의 trade-off관계 속에서 optimize를 시도하였으나 여전히 발전이 필요해 보인다. 

본 논문에서는 pixel domain이 아닌 frequency domain에서 low-frequency와 high-frequency로 분리함으로써  Deng의 부족한 점을 보완할 수 있었다. 새로운 SISR(single image super-resolution) method를 선보이는 것은 아니나, 기존의 성질이 다른 두 네트워크를 사용하여 얻은 두 결과 이미지를 fusion함으로써 best tradeoff를 찾았다는 것에 의의가 있다.  

본 논문을 SRGAN-MSE 논문과 비교하였을 때 deep network를 train하지 않아 stability 걱정이 없으며, ESRGAN과 비교하였을 때 네트워크 구조가 훨씬 flexible하고, Deng과 비교하였을 때 wavelet domain에서의 style transfer를 구현하여 새로운 테크닉을 선보였다는 점에서 좋은 결과를 얻을 수 있었다.

3 Proposed method

Stationary wavelet transform

기존의 discrete wavelet transform(DWT)는 shift invariant하지 않기 때문에 convolution을 적용할 수 없다. 따라서 DWT에서 downsampling operation을 제거한 stationary wavelet transform(SWT)을 적용하였다. 2D SWT는 그림3의 수식과 같이 1D wavelet decomposition인 H0(low-pass filter)과 G0(high-pass filter)의 z transform으로 얻어진다. 

2D SWT는 이미지를 여러개의 sub-bands로 나누는데, 이는 1개의 low-frequency sub-bands(LL)와 여러개의 high-frequency sub-bands(LH, HL, HH)로 구성된다. 2 level decomposition을 진행한 후, input X는 7개(3*2high+1low)의 상태로 나뉘어 지며 LH, HL, HH는 각각 horizontal, vertical, diagonal detail을 나타낸다. level2는 level1의 LL로부터 생성된다. 

Motivation

여기서 중요한 insight는 LL을 나타내는 low-frequency sub-band가 이미지의 objective quality에, 나머지 LH, HL, HH와 같은 high-frequency sub-bands가 perceptual quality에 각각 지대한 영향을 미친다는 것이다. 이를 증명하기 위해 CX로부터 얻은 high perceptual quality를 가진 Ap 이미지와 EDSR로부터 얻은 high objective quality를 가진 Ao 이미지를 얻었다. Ap, Ao 이미지를 1 level decomposion하여 얻은 sub-bands를 관찰해 보자. 

그림4를 보면, Ao의 LL sub-band histogram이 Ap의 LL보다 ground truth LL과 더욱 유사함을 볼 수 있고, Ap의 LH, HL, HH sub-bands histogram이 Ao의 LH, HL, HH보다 ground-truth의 LH, HL, HH와 더욱 분산이 유사함을 볼 수 있다. 

다음으로 정량적인 분석을 해보자. 성능비교를 위해 objective quality에는 PSNR(peak signal-to-noise-ratio)를, perceptual quality에는 NRQM을 각각 평가지표로 사용하였다. (PSNR과 NRQM은 모두 높을수록 좋은 성능임을 나타낸다.)

Ao의 LL과 Ap의 LH, HL, HH를 합친 것을 Ap~, Ap의 LL과 Ao의 LH, HL, HH를 합친 것을 Ao~라고 할 때, Ap, Ap~, Ao, Ao~의 PSRN, NRQM을 출력해 보았다. 표를 보기 전, 앞서 말했던 맥락을 참고하였을 때, Ap~는 Ao의 LL을 받았으므로 Ap보다 더 좋은 objective quality를 갖고 있어야 한다. 따라서 PSNR이 높아지고 NRQM이 조금 낮아질 것을 예상할 수 있다. Ao~는 Ap의 LL을 받았으므로 Ao보다 더 좋은 perceptual quality를 갖고 있어야 하나, LL은 다른 sub-bands에 비해 perceptual 정보를 적게 갖고 있으므로 NRQM이  조금 높아지고(Ap에 못미치게) PSNR은 낮아질 것을 예상할 수 있다.   

여기까지 보았을 때, 우리는 LLo가 objective quality를 높일 sub band이며, 나머지 sub band가 LHp, HLp, HHp에 가까울수록 perceptual qualiity가 높아질 것을 알 수 있다. 

Low-frequency sub-band enhancement (LSE)

이제 드디어 네트워크의 뼈대를 살표볼 차례이다. perceptuality를 높이는 네트워크(CX)에서 얻은 아웃풋(Ap)과 objectivity를 높이는 네트워크(EDSR)에서 얻은 아웃풋(Ao)을 각각 SWT(wavelet transform)하여 frequency domain으로 분리한다.

low frequency sub band인 LLo는 objective quality를 더욱 강화하기 위한 LSE 작업을 거친다. LSE는 VDSR의 구조를 차용한 작은 네트워크로, 6개의 convolution layer와 Relu를 지나 마지막에 input을 더해 output을 출력한다.  모든 layer의 filter 는 64개, filter size는 3x3이며 loss는 l2 norm을 이용한다. 

Wavelet domain Style transfer (WDST)

SWT 하여 얻은 LHo, HLo, HHo를 style transfer하여 LHp, HLp, HHp에 비슷하게 바꿈으로써 perceptual quality를 높이고자 한다. 따라서, LHp를 style input으로, LHo를 content input으로 하여 LHo의 content를 유지하며 LHp의 style을 갖는 아웃풋 LHr을 생성할 수 있도록 style transfer를 진행한다.

기존 style transfer와의 차이점은 wavelet coefficients를 input으로 넣는다는 것이다. wavelet transform한 결과는 음수도 있고 1보다 큰 숫자도 있기 때문에 VGG19 network에 input으로 넣기 전 0-1로 normalize해야한다.

normalization 후 각각의 high-frequency sub-band pair에 대하여 새롭게 정의한 loss를 줄이는 방향으로 아웃풋을 만들어 낸다. loss function은 수식1에서 볼 수 있듯 content loss(Lc), style loss(Ls), l1 loss의 weighted sum으로 구성된다. l1 norm loss를 추가한 이유는 wavelet coefficients에 0 값이 많기 때문이다(그림4의 히스토그램 참고). 데이터가 sparse할 때(0값이 많을 때), l2 loss보다 l1 loss를 사용하는 것이 더 효과적인데,  다음 사이트에 설명이 잘 나와있어 첨부한다. 간단히 말하자면 l2 loss는 항상 양수인 값들이 존재하기 때문에 sparse한 상황에서 0을 향해 optimize되는 데에 한계가 있다. 

https://stats.stackexchange.com/questions/45643/why-l1-norm-for-sparse-models

content loss와 style loss는 pre-trained VGG19 네트워크를 이용하여 여러 feature를 중간에서 뽑아 loss를 계산한다(그림5의 WDST). content loss는 MSE loss로 conv2-2 에서 LHo와 LHr의 feature를 한번 뽑아 content loss를 계산하며 수식2의 N은 layer L의 feature map 수, M 은 feature map의 weight*height를 나타낸다.

style loss는 Relu1-1, Relu2-1, Relu3-1, Relu4-1, Relu5-1 총 5개의 층에서 LHp와 LHr의 feature를 뽑아 weighted sum하여 style loss를 계산한다(수식3). style loss는 gram matix의 MSE loss를 계산함으로써 LHr과 LHp 간의 correlation을 계산한다. 

이렇게 high frequency sub-bands LHr, HLr, HHr를 얻었으면 de-normalization해주어 본래의 data 분포를 갖게끔 만들어준다. 앞서 LSE를 통해 강화된 LLr과 LHr, HLr, HHr를 합쳐 2D ISWT(inverse stationary wavelet transform)하면 결과 이미지를 얻을 수 있다. 

4 Numerical Results

Experimental setup

-2D SWT의 wavelet filter로는 bior2.2를 사용

-wavelet decomposition level은 2

-SGD optimizer 사용 (LSE)

-batch size 64 (LSE)

-basic learning rate 0.01, momentum 0.9 (LSE)

-ratio between the content loss and the style loss 10^(-3) (WDST)

-ratio between the content loss and the l1 norm loss 10^(-5) (WDST)

-style loss의 각 layer weight 0.2 (WDST)

-first layer maximum iteration 5000, second layer maximum iteration 1000 (WDST)

-Ao는 EDSR을 통해 얻음 

-Ap는 CX를 통해 얻음

-dataset : Set5, Set14, BSD100, Urban100, PIRM

Wavelet filter sensitivity

본 논문에서는 7가지의 wavelet filter를 사용하여 PSNR, SSIM, NRQM을 비교하였다. 다른 wavelet function들과는 다르게 rbior2와 bior4.4는 decomposition과 reconstruction시 사용되는 wavelet function의 모양이 다르다. PSNR을 높이는 데에는 haar, db2가 효과적이었으며, NRQM을 높이는 데에는 bior4.4가 효과적이었다. 

나는 여기까지 읽었을 때, "그래서 이렇게 하면 좋다는건 알겠는데, EDSR과 CX를 그냥 합친것보다 뛰어난지는 어떻게 알지?" 라는 의문이 들었다. 

Content and Style inputs sensitivity

그림8은 objectivity가 뛰어난 네트워크에서 뽑은 결과 Ao와 perceptuality가 뛰어난 네트워크에서 뽑은 결과 Ap의 단순 interpolation 결과와 본 논문의 방법을 비교한 것이다. PSNR과 Perceptual score모두 interpolation한 결과보다 뛰어난 성능을 보임을 알 수 있다. 

드디어 논문 리뷰가 끝이 났다!! 다음에는 cycleGAN을 다뤄볼까...

REFERENCE

1. Xin Deng et al, "Wavelet Domain Style Transfer for an Effective Perception-distortion Tradeoff in Single Image Super-Resolution", 2019 ICCV

2. l1 sparsity : https://stats.stackexchange.com/questions/45643/why-l1-norm-for-sparse-models

3. wavelet function : http://wavelets.pybytes.com/