본문 바로가기

AI 논문 정리/GAN (Generative Adversarial Network)

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

https://arxiv.org/abs/1711.09020

 

Conference on Computer Vision and Pattern Recognition (CVPR), 2018

 

1. Introduction


What is image-to-image translation??

주어진 이미지의 모습을 다른 모습으로 바꾸는 것을 의미한다. 두 개의 다른 도메인에서 학습 데이터가 주어졌을 때, StarGAN은 한 도메인에서 다른 하나의 도메인으로 바꾸는 것을 학습하는 모델이다.

 

CelebA 데이터셋과 같이, labeled 데이터셋을 사용하는 multi-domain image translation 과제를 수행할 때, 기존의 모델들이 비효율적인 것을 확인할 수 있다. k개의 domain 사이에서의 모든 매핑들을 학습하기 위해서, k(k-1)개의 generator가 학습되어야 하기 때문에 비효율적이다. 

즉, 위와 같이 4개의 도메인들 사이에서 image-to-image translation을 수행하기 위해서는 4 * 3 = 12개의 네트워크가 필요한 것이다. 각 데이터셋들은 부분적으로 labeling되어 있기 때문에 joint-training을 할 수 없다.

이러한 비효율성을 해결하기 위해서 연구진은 StarGAN을 제안하였다. StarGAN은 복수의 도메인들 사이의 mapping을 하나의 generator를 통해서 학습시킨다. 

 

Generator는 이미지와 도메인 정보를 input으로 받고, 알맞는 도메인으로 이미지를 translate한다. 도메인 정보를 나타내기 위해 label을 사용하는데, 주로 binary 또는 One-hot vector 형식으로 나타낸다. 

 

학습과정 중에, Target 도메인의 label을 무작위로 생성하고 input 이미지를 target 도메인으로 유연하게 translate할 수 있도록 모델을 학습시킨다. 이러한 과정을 통해서, 도메인 Label을 제어하는 것 뿐만 아니라, 이미지를 원하는 도메인으로 translate할 수 있게 된다.

 

또한, Mask Vector를 도메인 label에 추가하여, 기존의 모델에서는 안되던 joint-training을 할 수 있도록 하였다. 설명하자면, 모델이 모르는 label은 무시하고, 특정 데이터셋이 제공한 Label에만 집중할 수 있도록 해주는 방법이다. 

 

이 논문의 주요 contribution을 정리하면 다음과 같다.

1. 하나의 generator와 하나의 discriminator만을 사용해서 multiple domain 간의 매핑을 학습하는 StarGAN 제안함.

2. Mask Vector Method를 활용해서 다양한 데이터셋 간의 multi-domain image translation을 학습할 수 있음을 보여줌.

3. StarGAN의 질적, 양적인 결과를 제공해서 이전의 모델보다 좋은 성능임을 증명한다.

 

2. Related Work


Generative Adversarial Networks

 

GAN은 Image Generation, Image Translation, Super-Resolution, Face Image Synthesis와 같은 컴퓨터 비전 과제 수행에 좋은 성능을 보인 네트워크이다. 

보통 GAN 은 Discriminator와 Generator로 구성이 되어 있다. 

Discriminator는 진짜와 가짜 샘플들을 구분하도록 학습되고, Generator는 진짜 샘플과 구분이 되지 않을 정도로 정교한 가짜 샘플을 생성하도록 학습된다. 또한, 생성된 이미지들이 최대한 사실적으로 보이도록 adversarial loss를 조절한다. 

 

Conditional GANs

 

이전 연구에서는 Discriminator와 Generator 모두에게 class에 대한 정보를 제공해서 class를 기반으로 하는 샘플들을 생성하도록 유도하였다. 최근 연구에서는 문자 설명을 기반으로 특정한 이미지를 생성하도록 유도하기도 한다.

 

Image-to-Image Translation

 

pix2pix가 cGAN을 사용해서 좋은 성과를 내었다. 또한, CycleGAN과 DiscoGAN도 괄목할만한 성과를 내었다. 하지만, multiple-domain을 다룰 때에는 한계가 있으며, StarGAN은 이를 보완한 모델이다. 

 

 

3. Star Generative Adversarial Networks (StarGAN)


StarGAN은 하나의 데이터 셋 내에서 multi-domain image-to-image translation을 하는 프레임워크이다. 먼저, multi-domain image-to-image translation이 무엇인지부터 설명하도록 하겠다.

StarGAN

 

3.1 Multi-Domain Image-to-Image Translation

 

이 논문은 여러 도메인간의 Mapping을 배우는 하나의 Generator G를 학습시키는 것이 목표이다. 

 

이를 위해, input 이미지 xtarget domain label c를 기반으로 output 이미지 y로 변환시킬 수 있도록 G를 학습시키게 된다.

(c는 무작위로 생성되어 G가 input 이미지를 유연하게 변환시킬 수 있도록 한다)

또한, 이 논문은 하나의 discriminator가 여러 도메인을 조절하기 위해서 auxiliary classifier를 도입한다. 이를 통해 Discriminator는 source와 domain label에 대한 확률분포를 생성하게 된다. 

(**auxiliary classifier -> Classification의 문제를 해결하는 Neural Network는 softmax를 맨 마지막 layer에 딱 하나만 놓는데, Auxiliary classifier는 중간중간에 softmax를 두어 중간에서도 Backpropagation을 하게 함. 결론적으로 Vanishing Gradient 문제를 해결하도록 돕는다)

 

지금부터는 StarGAN의 loss function에 대해서 정리해보았다.

 

Adversarial Loss

 

새로 생성된 이미지들이 실제 이미지들과 구분되지 않도록 하기 위해서, Adversarial Loss를 도입한다. 

G는 이미지 G(x,c)를 생성해내고, D는 진짜 이미지와 가짜 이미지를 구분하게 된다. D_src(x)는 D에 의해서 주어진 source에 대한 확률 분포이다. G는 loss 값을 최소화하려는 반면에, D는 loss 값을 최대화하려고 한다. 

 

Domain Classification Loss

앞서 언급한 것과 같이, StarGAN은 input 이미지 x를 target domain label c를 기반으로 x를 output 이미지 y로 변환시키는 것이 목표이다. 이를 제대로 수행하기 위해서, auxiliary classifier를 discriminator 위에 추가하여 domain을 분류한다.

여기에서 loss를 두 파트로 나누게 된다.

 

1) Domain Classification Loss of Real Images used to optimize G

D_cls(c'|x)는 D가 compute한 도메인 label에 대한 확률분포이다. 이 loss를 최소화시킴으로써, D는 진짜 이미지 x를 원래의 도메인 c'으로 분류시키는 방법을 학습하게 된다. (x 와 c'의 값은 training data가 제공해준다고 가정한다.)

 

2) Domain Classification Loss of Fake Images used to optimize G

G는 이 loss를 최소화시킴으로써 target domain c로 분류될 수 있는 이미지를 생성시킨다.

 

 

Reconstruction Loss

앞서 정리한 Adversarial Loss와 Classification Loss를 최소화시킴으로써, G는 target domain에 제대로 분류되는 사실적인 이미지를 생산하도록 학습이 되어진다. 

하지만, loss를 최소화시키는 것이 변환된 이미지가 input 이미지의 내용을 보존한다는 것을 보장하지는 않는다. 이러한 문제를 해결하기 위해서, 이 논문은 Cycle Consistency Loss를 도입한다. 

G는 변환된 이미지 G(x, c)와 오리지널 domain label c'를 input으로 받고 오리지널 이미지 x를 다시 생성하게 된다. 이를 Reconstruction 단계라고 부른다. 

 

L1 norm이 reconstruction loss를 사용되어졌다. 유의해야 할 점은, Single Generator를 두 번 사용했다는 점이다. (Original -> Image in target domain & translated image -> reconstruct original image)

 

 

Full Objective 

위에서 정리한 loss function들을 정리하면 다음과 같은 식으로 정리할 수 있다.

실험에서는 λ_cls = 1,  λ_rec = 10으로 한다.

 

 

 

3.2 Training with Multiple Datasets

 

StarGAN의 굉장히 중요한 장점은 다른 label들을 가지고 있는 다양한 데이터 셋을 동시에 처리할 수 있다는 점이다. 

 

하지만, 이와 같이 복수의 데이터 셋에서 학습을 할 때, label의 정보가 부분적으로만 알려져 있다. 예를 들어서, CelebA 데이터셋에서 Hair Color와 Gender와 같은 label은 있지만, happy와 angry와 같은 label을 가지고 있지 않다.

Reconstruction을 할 때, 모든 label 정보에 대한 label vector c'가 필요하기 때문에 이 문제가 발생하는 것이다. 

 

Mask Vector

부분적으로만 알려져 있는 label 문제를 해결하기 위해서, 논문은 Mask Vector m이라는 개념을 도입하였다. Mask Vector는 StarGAN이 불분명한 label을 무시하고 명확하게 알려진 label에 대해서만 집중할 수 있도록 해준다.

 

StarGAN에서는 n-차원의 one-hot vector를 사용해서 m을 나타낸다. (n = 데이터 셋의 개수)

 

Label의 통합된 버전은 다음과 같다. 

c_i의 vector는 binary vector 또는 one-hot vector형식으로 표현된다. 나머지 n-1의 label들에 대해서는 0의 값을 배정한다.

 

이 논문에서 사용한 CelebA와 RaFD 데이터셋은 n = 2이다.

 

 

Training Strategy

StarGAN은 domain label

를 Generator의 input으로 사용한다. 이를 Generator의 인풋으로 사용하면, 불분명한 label (zero vector)는 무시하고, 명확하게 알려진 label만 집중할 수 있게 된다. 방금 언급한 input label을 제외하고는 single dataset을 사용해서 학습을 할 때와 Generator의 구조는 동일하다. 

 

Discriminator의 Auxiliary Classifier의 범위를 넓혀서 모든 데이터 셋의 Label에 대한 확률분포를 생성한다. 그 이후, multi-task 학습 환경에서 모델을 학습하고, Discriminator는 알고 있는 label에 대한 classification error만을 최소화시키고자 노력한다.

 

다양한 데이터 셋으로 학습시키면서, Discriminator는 모든 데이터 셋의 discriminative feature를 학습하고, Generator는 모든 데이터 셋의 label을 제어하는 것을 학습한다. 

 

 

4. Implementation


Improved GAN Training

학습과정을 안정화하고, 좋은 화질의 이미지를 생성해내기 위해서, Adversarial Loss의 수식을 Gradient penalty와 함께 Wasserstein GAN 수식으로 대체해서 사용하였다. 

 

Network Architecture

StarGAN은 CycleGAN을 기반으로 만들어진 모델이며, 두 개의 convolutional layer들 (Stride size = 2 for downsampling, 6 residual blocks & two transposed convolutional layers with stride size = 2 for upsampling) 로 구성이 된 Generator Network를 가진다.

Generator에서는 instance normalization을 사용하지만, discriminator에서는 normalization을 사용하지 않는다.

또한, Discriminator를 위해서 PatchGAN의 영향력을 조절하는데, 이를 통해 이미지가 실제인지 가짜인지 분류하게 된다.

 

 

6. Conclusion


결론적으로, StarGAN은 하나의 Generator와 하나의 Discriminator를 사용해서 multiple domain에 걸쳐 Image-to-image translation이 가능해지도록 한다. StarGAN은 확장성이 좋을 뿐만 아니라, 기존의 연구에 비해서 뛰어난 visual quality를 가진 이미지를 생성하였다.

또한, 이 논문에서 제안한 Mask Vector를 통해서 StarGAN이 다른 label을 가진 다수의 데이터 셋을 동시에 다룰 수 있도록 하였다. 

 

 

참조 문서: https://technical-support.tistory.com/87 - Auxiliary Classifer 설명