[Paper Review] StarGAN : Unified Generative Adversarial Networks for Multi-Domain Image-to-Image Translation

Main Reference:
- CycleGAN Paper

🚀 Introduction

Image-to-Image translation for two domains 작업은 이전에 리뷰한 pix2pix, CycleGAN 등 좋은 성과를 내고있다. 하지만 기존 모델들은 domain pair마다 별도의 Generator & Discrimator를 훈련시켜야 하기 때문에 scalability&robustness에 한계가 존재한다. 다시 말해 k개의 multi-domain을 학습하려면 k(k-1)개의 Generator, Discriminator가 필요한 것이다. 한편 서로 다른 domain들이 공통적으로 갖는 feature가 존재함에도, 각각의 dataset들은 부분적으로 라벨링 되어있다. 이러한 partially labeled 된 dataset들을 jointly training 할 수 없기 때문에 비효율적인 학습 방법이라 말할 수 있다. 본 논문에서는 앞서 언급한 문제들을 해결하기 위해 하나의 Generator로 multi-domain을 학습할 수 있는 StarGAN 모델을 제안한다.

🚀 Star Generative Adversarial Networks

StarGAN의 목적은 하나의 Generator가 multiple-domain 사이의 매핑을 학습하는 것이다. 이를 위해 Discriinator에 Auxiliary Classifier를 추가했고, target domain label을 정해 conditional로 Generator를 학습했다.

Objective Function

\[\textbf{Adversarial Loss : } \quad \textbf{L}_{adv} = \mathbb{E}_{x}[logD_{src}(x)] + \mathbb{E}_{x,c}[log(1 - D_{src}(G(x,c)))]\] \[\textbf{Domain Classification Loss : } \quad \textbf{L}_{cls}^{r} = \mathbb{E}_{x,c'}[ - log(D_{cls}(c'|x))]\] \[\textbf{Reconstruction Loss : } \quad \textbf{L}_{rec} = \mathbb{E}_{x,c,c'}[\left\| x - G(G(x, c), c') \right\|_{1}]\] \[\textbf{L}_{D} = - \textbf{L}_{adv} + \lambda_{cls} \textbf{L}_{cls}^{r}\] \[\textbf{L}_{G} = \textbf{L}_{adv} + \lambda_{cls} \textbf{L}_{cls}^{f} + \lambda_{rec} \textbf{L}_{rec}\]

• Discriminator, Generator Loss에 모두 classification loss를 반영
• c : labels of target domain / c’ : labels of original domain
• Adversarial loss나 Classification loss만으로는 바꾼 label 이외의 content들이 보존되는 것을 보장하지 않는다.
  • 따라서 Generator loss에 cycle consistency loss(Reconstruction Loss)를 추가
• 람다값들은 dataset마다 다른데 저자들은 cls:1, rec:10 사용

Mask Vector

Multiple datasets을 학습할 때 발생하는 이슈는 label 정보가 dataset의 일부 정보만을 포함한다는 것이다. 예를 들어 celeb dataset에 포함되는 이미지는 표정에 대한 label을 가지고 있지 않는 상황이다. 이러한 문제들을 해결하기 위해 unspecified label을 무시하고, 알려진 특정 label들에 집중하게 해주는 mask vector를 도입했다. One-hot vector 아이디어를 이용한 단순한 개념이다.

Training Strategy

🚀 Implementation

Improved GAN Training

\[\textbf{L}_{adv} = \mathbb{E}_{x}[logD_{src}(x)] + \mathbb{E}_{x,c}[D_{src}(G(x,c))] - \lambda_{gp} \mathbb{E}_{\hat{x}}[(\left\| \bigtriangledown_{\hat{x}} D_{src}(\hat{x}) \right\|_{2} - 1 )^{2}]\]

안정적인 학습을 위해 gradient panalty를 이용한 Wasserstein GAN loss를 사용하였다(lambda=10).

Network Architecture

🚀 Experiments

Celeb A Dataset	AMT Score

• 단일 Dataset 내에서도 StarGAN의 AMT socre가 가장 높음 (정량적 평가)

• SNG(single training)보다 JNT(joint training) 방법이 더 high quality 이미지 생성

🚀 Code

StarGAN from scartch (pytorch)

Original Hair	Black Hair

맨 위로 이동하기

Paper 카테고리 내 다른 글 보러가기