노션으로 보시면 더 편할 수 있습니다.
[논문 정리] AdaIN (Arbitrary Style Transfer in Real-Time With Adaptive Instance Normalization)
어떤 사람이 이 글을 보는게 좋은가?
www.notion.so
어떤 사람이 이 글을 보는게 좋은가?
- 기본적인 ML지식을 갖추고 있으신 분
- Style Transfer에 관심이 있으나, 이제 막 시작해서 알고싶으신 분
- AdaIN의 수식에 대해서 궁금하신 분
- AdaIN 소스코드의 간단한 리뷰가 궁금하신 분
요약
최초의 Neural Style Transfer를 제안한 Gatys의 Style Transfer 방식은 다양한 Style을 Arbitrary하게(그 때 그 때 새로운 스타일을) 적용할 수 있는 반면에, 굉장히 느린 속도로 style transfer를 수행한다는 단점이 있었다.
이를 극복하기 위해 Feed-forward 방식으로 Style Transfer를 수행하는 방식들이 제안되었는데, 이들은 Gatys 방식의 비해서 빠른속도의 Style Transfer가 가능했으나, 한정적인 몇가지 미리 학습된 Style 에대해서만 Style Transfer가 가능했다.
이에 비해 AdaIN 방식은 빠른속도로 추론이 가능하면서 동시에 Arbitrary하게 새로운 스타일을 적용할 수 있는 방식이다.
위 표의 Method를 설명해보자면 Gatys - 최초 논문, Ulyanov - Instance Normalization(IN), Dumoulin - Conditional Instance Normalization(CIN), our - Adaptive Instance Normalization (AdaIN) 이다.
실험 결과를 보면 Gatys는 많은 스타일에 대해서 적용할 수 있지만, 속도가 굉장히 느린 것을 확인할 수 있다. 이에 비하여 Feed-forward방식의 IN 과 CIN은 빠른속도의 추론이 가능했지만, 스타일이 한정적이다. AdaIN은 이러한 방식들과는 차별화되게 빠른 속도의 추론이 가능하면서 동시에 무한한 스타일을 생성해 낼 수 있다는 장점을 가지고 있다.
Architecture
AdaIN의 네트워크 구조는 아래와 같으며, 인상 깊은 점은 녹색의 VGG의 pre-trained 모델을 통해서 Encoding을 수행하며, 이 encoder를 feature를 인코딩할 때, 그리고 Loss Function을 구할 때 사용한다는 것이다. 즉 Encoder는 학습 시키지 않는다는 점이 포인트다.
그러므로, 자연스럽게 이 네트워크 상에서 학습 시키는 것은 Decoder 뿐이며, 저자들의 표현을 빌리자면 이 Decoder는 AdaIN으로 생성된 feature들이 decoder를 통해서 image space로 invert 하는 법을 학습한다. 아직 설명하지 않았지만, AdaIN 내에서는 learnable parameter가 없다.
AdaIN Layer
그렇다면 AdaIN은 어떻게 생겼는가? 이를 알기 앞서서 Style Transfer의 개념에 대해서 간단하게 알고 있어야 한다. Style Transfer는 특정 이미지에서 Style을 뽑고, 다른 이미지에서 Contents를 뽑아서 이를 합성한다.
Style은 직관적으로 와닿는데 Contents는 무엇인지 잘 이해가 안갈 것이다. 간단하게 어떤 형태라고 생각하면 좋다. 나무의자가 있다면, 의자가 Contents, 나무가 Style이 될 것이다.
AdaIN에서는 Style과 Contents에 대한 정보를 VGG Encoder를 통해서 추출할 수 있다고 주장한다 사실 이러한 아이디어는 초기 연구인 Gatys 때부터 사용해왔었다.
아주 조금만 더 깊게 들어가서...초기 연구인 Gatys 방식에서는 VGG에서 나온 Feature들에 Gram Matrix를 사용해서 Style을 표현(representation) 있다는 것을 실험적으로 보였다. 이러한 Gram Matrix가 대표적인 feature space 상의 statistics를 추출해내는 방법인데, 이 이후 많은 연구가 이루어지면서 feature space상의 여러 statistics가 Style을 표현하는데 유용하다는 것이 실험적으로 많이 밝혀졌다.
statistics? 잘모르겠다고 생각한다면 이 논문에서는 그냥 평균(mean)과 분산(variance)라고 생각해도 무방하다. AdaIN은 feature space 상의 평균과 분산이 Style에 영향을 끼친다면, 이들을 뽑아서 즉석으로 교환해주는 방식을 택한 것이다. 식을 보자면 아래와 같다.
$$\operatorname{AdaIN}(x, y)=\sigma(y)\left(\frac{x-\mu(x)}{\sigma(x)}\right)+\mu(y)$$
이 식에서 $\mu$ 함수는 평균을 구하는 함수이고, $\sigma$ 함수는 표준 편차를 구하는 함수이다. 자주 사용되는 term이다.
Style Transfer의 경우, 내가 원하는 Contents를 담고 있는 이미지의 feature $x$ 에서, 이미지의 스타일을 빼주고, 내가 입히고 싶은 Style을 더해주는 방식으로 수행된다.
그런데, 위에서도 말했 듯 스타일은 feature 상에서의 statistics로 표현된다고 말했는데 그게 바로 평균과 분산이다. 따라서 식에서 보자면 $\left(\frac{x-\mu(x)}{\sigma(x)}\right)$ 는 Contents 이미지에서 Contents 이미지의 스타일을 빼준 것이고, $\sigma(y)\left(\frac{x-\mu(x)}{\sigma(x)}\right)+\mu(y)$ 는 이미지 y의 스타일을 입혀준 것이다. 이해가 됐으려나...
그리고 했깔릴 수 있는 부분인데, 이건 전부 Feature space상에서 이루어진다는 것이다! 그냥 이미지상에서 이루어지는 것이 아니다.
Formulation
노테이션을 정리해보자면
- $T$= Style Transfer Network (인코더-AdaIN-Decoder)
- $f$= Encoder (pre-trained VGG-19의 앞부분 (~relu4_1)
- $g$= Decoder (학습시켜야하는 디코터)
라고 할 때, AdaIN layer를 통해서 생성되는 feature $t$ 는 아래와 같이 나타낼 수 있다.
$$t=\operatorname{AdaIN}(f(c), f(s))$$
랜덤하게 초기화된 디코더 $g$는 $t$를 image space로 보내는 방법을 트레인하며 스타일이 입혀진 이미지 $T(c,s)$를 생성한다.
$$T(c, s)=g(t)$$
Architecture Detail
- checker-board effect를 감소시키기 위하여 decoder의 pooling layer를 nearest up-sampling 방식으로 교체.
- $f$와 $g$에서 모두 reflection padding을 사용했다.
- decoder에서 normalization 방식을 고르는 것이 중요했는데, 결론은 no normalization이 제일 좋았다.
- 전처리: 두 이미지를 aspect ratio를 유지한체 512로 사이즈를 키웠고, 여기서 256 by 256로 crop한다. 우리의 네트워크는 fully convolutional 이기 때문에, 어떤 사이즈의 이미지가 온다고해도 적용 가능하다.
Loss
Loss는 다른 네트워크와 유사하게 다음과같은 로스를 사용한다.
$$\mathcal{L}=\mathcal{L}_{c}+\lambda \mathcal{L}_{s}$$
content loss는 target feature와 output image의 feature의 Euclidean distance를 구했다. 일반적으로 사용되는 content image의 feature response를 사용하는 대신에 AdaIN output $t$를 content target으로 삼았다. 이게 조금더 빨리 convergence가 이루어진다.
$$\mathcal{L}_{c}=\parallel f(g(t))-t \parallel_{2}$$
그냥 전체 아키텍처 오버뷰의 보라색 화살표의 식이다. $t$를 디코더에 넣고 다시 인코더에 넣은 후에 그 두개의 차이를 비교하는 것을 Contents Loss로 삼겠다는 간단한 식이다.
AdaIN 레이어는 오직 style features의 mean과 standard deviation를 transfer하기 때문에 style loss는 이러한 statistics를 match시켜야한다. 따라서 아래와 같은 스타일 로스를 사용하는데,
$$\begin{array}{r}\mathcal{L}_{s}=\sum_{i=1}^{L}\left\|\mu\left(\phi_{i}(g(t))\right)-\mu\left(\phi_{i}(s)\right)\right\|_{2}+ \\\sum_{i=1}^{L}\left\|\sigma\left(\phi_{i}(g(t))\right)-\sigma\left(\phi_{i}(s)\right)\right\|_{2}\end{array}$$
$\phi_i$.는 VGG-19의 i번째 레이어이다. 스타일로스에서 사용한 레이어는 relu1 1, relu2 1, relu3 1, relu4 1 이다. 이 역시 간단하게 설명하자면, 원래 스타일 $s$ 를 인코더에 넣었을 때의 $i$ 번째 feature $\phi_i(s)$ 의 평균과 $t$ 를 디코더에 넣고 이를 다시 encoder에 넣었을 때의 $i$번째 Feature $\phi_i(g(t))$ 의 평균($\mu$)과 표준편차($\sigma$)를 최소화 시키는 방법으로 스타일 로스를 구한것이다.
Code
사실 이렇게 길게 설명했는데, 코드로는 아래와같이 씸플하다.
def adaptive_instance_normalization(content_feat, style_feat):
assert (content_feat.size()[:2] == style_feat.size()[:2])
size = content_feat.size()
style_mean, style_std = calc_mean_std(style_feat)
content_mean, content_std = calc_mean_std(content_feat)
normalized_feat = (content_feat - content_mean.expand(
size)) / content_std.expand(size)
return normalized_feat * style_std.expand(size) + style_mean.expand(size)AdaIN의 코드는 그냥 content_feature와 style feature를 수식 그대로 적용해 반환할 뿐인 함수다.
def forward(self, content, style, alpha=1.0):
assert 0 <= alpha <= 1
style_feats = self.encode_with_intermediate(style)
content_feat = self.encode(content)
t = adain(content_feat, style_feats[-1])
t = alpha * t + (1 - alpha) * content_feat
g_t = self.decoder(t)
g_t_feats = self.encode_with_intermediate(g_t)
loss_c = self.calc_content_loss(g_t_feats[-1], t)
loss_s = self.calc_style_loss(g_t_feats[0], style_feats[0])
for i in range(1, 4):
loss_s += self.calc_style_loss(g_t_feats[i], style_feats[i])
return loss_c, loss_sencode_with_intermediate 함수는 중간중간 레이어를 추출해서 적용하는 함수고, encode 함수는 그야말로 vgg를 relu4_1까지 통화시킨 후 feature map을 반환하는 함수다.
주목할만한 점은 adain에 style_feature의 마지막 것만 들어가는건데, 생각해보면 당연하다. shape이 마지막껏만 content feature와 맞기 때문이다. 그리고 이래도 style의 mean과 variance를 adaIN을 통해 transfer할 수 있다.
그리고 생성된 $t$ 를 디코더를 통해 이미지로 만들고, 이를 다시 encoder에 집어넣어서 contents loss와 style loss를 구해서 반환해준다.
논문의 전체 번역본과 소스코드 주소는 아래와 같다.
[번역] Arbitrary Style Transfer in Real-Time With Adaptive Instance Normalization