작성자 : 건국대학교 응용통계학과 정재윤

Image to Image Translation

image to image translation이란 이미지 데이터셋을 사용하여 input 이미지와 output 이미지를 mapping하는 것을 목표로 하는 생성모델의 한 분야입니다. 흑백 이미지에 컬러를 입힌다든지, 낮 사진을 밤 사진으로 만든다든지, 테두리만 주어진 사진을 실제 물건같이 만드는 것이 가능하죠.

이런 image to image translation은 크게 paired한 데이터셋을 사용하는 분야와 unpaired한 데이터셋을 사용하는 분야로 구분할 수 있습니다. 그리고 이 분야들을 대표하는 모델들이 오늘 볼 pix2pix과 cyclegan입니다.

Pix2Pix

그 당시 이미지 처리 문제의 대부분은 CNN을 사용해서 해결했습니다. 하지만 그 당시 CNN에는 큰 결점이 하나 있습니다. 특정 과업을 수행하기 위해서는 매번 특정 loss를 설계해야하기 때문에 포괄적인 작업을 할 수 없다는 것입니다. 실제 pix2pix에서 진행하는 실험들을 CNN으로 시도한 적 있으나 대부분의 결과는 blurry했습니다. 그래서 논문의 저자들은 CGAN이라는 조건부 생성모델을 바탕으로 포괄적인 작업을 수행할 수 있는 pix2pix을 만들었습니다.
CGAN은 지난 번 민정님께서 잘 설명해주셨으니 넘어갈게요 :)

기본적으로 pix2pix은 생성모델이기 때문에 generator, discriminator가 필요합니다.

generator의 경우, CGAN과 굉장히 유사하게 작동합니다. 하지만 input으로 Noise vector와 class vector를 받는 CGAN과는 다르게 더 이상 Noise vector를 필요로 하지 않습니다. 오로지 이미지만을 input으로 받아 output을 생성하는 것입니다.

discriminator의 경우, 2가지를 input으로 받습니다. 기본적으로 generator에 넣었던 input 이미지, 이와 하나의 pair가 될 진짜 이미지나 생성된 이미지를 input으로 받습니다. 이 2가지를 동시에 discriminator에 넣었을 때, discriminator는 pair를 비교하여 생성된 이미지인지 진짜 이미지인지를 구분하는 것입니다. 이 일련의 과정으로 학습을 마치게 됩니다.

여기서 알아두셔야 할 점은 기존의 discriminator와 generator를 한 단계 발전시켰다는 점입니다. generator는 U-net구조와 skip connection을 사용했고, discriminator는 patchgan을 사용했다는 점입니다.

U - net

[출처 : https://taeoh-kim.github.io/blog/image2image/]

기본적으로 U-net은 encoder-decoder 구조를 skip-connection으로 발전시킨 구조입니다. 데이터가 인코더를 통하면 이미지의 feature들을 잘 추출해낼 수 있겠죠? 그리고 다시 디코더를 통하면 원하는 이미지를 만들 수 있을 것입니다. 그러나 이 과정에서 이미지의 형태나 해상도 등 정보의 손실은 어쩔 수 없이 발생하게 됩니다. 이러한 정보의 손실을 막기 위해서 skip - connection을 사용합니다. 이 뿐만 아니라 skip-connection을 사용하면 역전파 과정에서 vanishing gradient 문제도 해결할 수 있다는 장점이 있습니다.

[출처 : https://taeoh-kim.github.io/blog/image2image/]

모델의 구조는 위와 같습니다. 각각 8개의 block으로 구성되어 있으며, 유의하셔야 점은 decoder의 처음 3개의 block에서 Batch Normalizaion과 Dropout이 같이 진행되는 것입니다.

    def __init__(self, image_channels=3, inner_channels=64, n_layers=8, dropout=0.5, norm_layer=nn.InstanceNorm2d):
        super().__init__()
        assert n_layers >= 5

        block = UNetSkipConnectionBlock(inner_channels*8, inner_channels*8, 'innermost', norm_layer=norm_layer)
        for _ in range(n_layers-5):
            block = UNetSkipConnectionBlock(inner_channels*8, inner_channels*8, 'middle', block, dropout, norm_layer=norm_layer)
        block = UNetSkipConnectionBlock(inner_channels*4, inner_channels*8, 'middle', block, dropout, norm_layer=norm_layer)
        block = UNetSkipConnectionBlock(inner_channels*2, inner_channels*4, 'middle', block, dropout, norm_layer=norm_layer)
        block = UNetSkipConnectionBlock(inner_channels, inner_channels*2, 'middle', block, dropout, norm_layer=norm_layer)
        self.model = UNetSkipConnectionBlock(image_channels, inner_channels, 'outermost', block)

    def forward(self, x):
        return self.model(x)

generator 구조 코드 [출처 : https://www.secmem.org/blog/2020/07/19/pix2pix/]

    def __init__(self,
                 outer_channels,
                 inner_channels,
                 module_type,
                 submodule=None,
                 dropout=0.5,
                 norm_layer=nn.InstanceNorm2d
                 ):
        super().__init__()
        
        if module_type not in ['innermost', 'outermost', 'middle']:
            raise Exception('no such module type')

        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d
        
        down_conv = nn.Conv2d(outer_channels, inner_channels, kernel_size=4, stride=2, padding=1, bias=use_bias)
        down_relu = nn.LeakyReLU(0.2, True)
        down_norm = norm_layer(inner_channels)

        up_relu = nn.ReLU(True)
        up_norm = norm_layer(outer_channels)

        self.outermost = module_type == 'outermost'
        if module_type == 'innermost':
            up_conv = nn.ConvTranspose2d(inner_channels, outer_channels, kernel_size=4, stride=2, padding=1, bias=use_bias)
            modules = [down_relu, down_conv, up_relu, up_conv, up_norm]
        elif module_type == 'outermost':
            up_conv = nn.ConvTranspose2d(inner_channels*2, outer_channels, kernel_size=4, stride=2, padding=1)
            modules = [down_conv, submodule, up_relu, up_conv, nn.Tanh()]
        else:
            up_conv = nn.ConvTranspose2d(inner_channels*2, outer_channels, kernel_size=4, stride=2, padding=1, bias=use_bias)
            modules = [down_relu, down_conv, down_norm, submodule, up_relu, up_conv, up_norm, nn.Dropout(dropout)]

        self. model = nn.Sequential(*modules)

    def forward(self, x):
        if self.outermost:
            return self.model(x)
        return torch.cat([x, self.model(x)], 1)

skip-connection 구조 코드 [출처 : https://www.secmem.org/blog/2020/07/19/pix2pix/]

PatchGAN

pix2pix의 discriminator는 PatchGAN입니다. PatchGan이라고 대단한 구조는 아닙니다. 기존의 GAN의 discriminator는 Image 전체를 보고 진짜인지 아닌지를 판별한다면, PatchGAN은 이미지를 조각내서 부분부분을 보고 진짜인지 아닌지를 판별합니다. 아래와 같은 그림처럼 말이죠. 그러면 Patch 단위로 loss의 역전파가 진행되기 때문에 좀 더 detail한 image를 만들 수 있게 되는 것이죠.

이 때 대체로 patch의 크기는 70 x 70 사이즈를 사용합니다. 그 이유는 이미지의 크기와 동일한 286 x 286 사이즈의 patch를 사용해도 70 x 70 사이즈의 결과물과 다를 바가 없고, 오히려 parameter만 증가해 학습시간에 악영향을 미치기 때문이라고 합니다. 모델의 구조는 아래와 같습니다.

[출처 : https://taeoh-kim.github.io/blog/image2image/]

    def __init__(self,
                 input_channels=6,
                 inner_channels=64,
                 n_layers=3,
                 norm_layer=nn.InstanceNorm2d
                 ):
        super().__init__()

        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d

        modules = [nn.Conv2d(input_channels, inner_channels, kernel_size=4, stride=2, padding=1), nn.LeakyReLU(0.2, True)]
        for i in range(n_layers-1):
            modules += [
                nn.Conv2d(inner_channels*min(2**i, 8), inner_channels*min(2**(i+1), 8), kernel_size=4, stride=2, padding=1, bias=use_bias),
                norm_layer(inner_channels*min(2**(i+1), 8)),
                nn.LeakyReLU(0.2, True)
            ]
        modules += [
            nn.Conv2d(inner_channels * min(2 ** (n_layers-1), 8), inner_channels * min(2 ** n_layers, 8), kernel_size=4, stride=1,
                      padding=1, bias=use_bias),
            norm_layer(inner_channels * min(2 ** n_layers, 8)),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(inner_channels * min(2 ** n_layers, 8), 1, kernel_size=4, stride=1, padding=1)
        ]
        self.model = nn.Sequential(*modules)

    def forward(self, x):
        return self.model(x)

discriminator 구조 코드 [출처 : https://www.secmem.org/blog/2020/07/19/pix2pix/]

Loss

pix2pix의 기본적인 loss는 위와 같은 형식입니다. 굉장히 익숙한 앞부분과 생소한 뒷부분으로 구성되어 있죠.

pix2pix은 생성모델이므로 adversarial loss를 사용합니다. 이 때, loss는 주로 BCE loss를 사용한다고 합니다.

위의 adversarial loss만을 사용하면 생성된 이미지가 약간 흐리게 나온다고 합니다. 그래서 논문의 저자들은 조금이라도 개선된 이미지를 만들기 위해 L1 distance (Manhattan distance)를 사용합니다. 즉, 만들어진 이미지와 실제 이미지 사이의 픽셀간의 거리를 구해 이를 최소화하여 최대한 원래 이미지처럼 만들 수 있게 하는 것입니다.

그래서 최종적으로 제일 위의 loss function을 만들고, 람다라는 hyper parameter를 조절하여 최적의 loss function을 찾아갑니다.

CycleGAN

위에서 설명한 pix2pix 모델은 paired한 이미지 데이터셋을 가지고 학습을 진행했습니다. 하지만 paired 데이터셋은 실제로 구하기 굉장히 어렵고, 구한다고 하더라도 데이터셋에 들어가는 비용이 만만치 않습니다. 뿐만 아니라 아예 못구하는 데이터들도 수두룩합니다. 그래서 논문의 저자들은 짝지어진 예시 없이 X라는 도메인으로부터 얻은 이미지를 타겟 도메인 Y로 바꾸는 방법인 CycleGAN을 제시합니다.

paired dataset의 경우, X 이미지의 좌표값과 Y 이미지의 좌표값이 어떻게 대응되는지에 대한 정보들이 담겨있습니다. 그러나 unpaired dataset은 일대일로 매칭된 것이 아니기 때문에 이에 대한 정보가 존재하지 않습니다. 이를 개선하고자 저자들은 cycle consistency loss function을 사용한 것입니다.

구조

기본적으로 CycleGAN은 2개의 GAN을 필요로 합니다. X에서 Y의 이미지를 만들어주는 generator와 이 이미지가 진짜인지 판단하는 discriminator, 그리고 반대의 경우까지 고려한 것이죠.

기본적으로 discriminator는 pix2pix과 마찬가지로 patchGAN을 사용합니다. generator의 경우, pix2pix과 같은 U - net을 기반으로 발전시켰다. 아래와 같이 U-net에 DCGAN과 residaul connection을 추가한 방식을 사용했습니다. 이렇게 만들면서 확실히 정보의 손실량이 줄어, 고해상도 처리를 더 수월하게 만들었습니다.

    if netG == 'resnet_9blocks':
        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
    elif netG == 'resnet_6blocks':
                net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
    elif netG == 'unet_128':
        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
    elif netG == 'unet_256':
        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
    else:
        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
    return init_net(net, init_type, init_gain, gpu_ids)

일반적으로 작동원리는 아래의 그림과 같습니다. 실제 이미지가 generator를 통하면 가짜 이미지를 생성합니다. 이 때, discriminator는 이 이미지가 진짜인지 가짜인지를 자신이 학습한 데이터를 토대로 판별하게 됩니다. output은 patchGAN을 사용했기 때문에 그림과 같이 각 부분마다 진짜인지 아닌지를 판별해서 구해집니다.

    # define networks (both Generators and discriminators)
    # The naming is different from those used in the paper.
    # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
    self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                    not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
    self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                    not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)

    if self.isTrain:  # define discriminators
        self.netD_A = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
                                        opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
        self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
                                        opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)

---

def forward(self):
    """Run forward pass; called by both functions <optimize_parameters> and <test>."""
    self.fake_B = self.netG_A(self.real_A)  # G_A(A)
    self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
    self.fake_A = self.netG_B(self.real_B)  # G_B(B)
    self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))

Loss

기본적으로 논문에서 제시한 loss function은 처음 이미지와 같습니다. GAN이 2개이기 때문에 각각의 adversarial loss와 cycle consistency loss를 고려한 것이죠. 이 후 이 function을 논문에서 차차 발전시켜서 두 번째 이미지와 같은 loss function을 만들게 됩니다. 이를 하나하나 살펴보겠습니다.

Cycle Consistency loss

Cycle Consistency를 사용하는 이유는 저희가 unpaired dataset을 사용하기 때문입니다. 위에서 언급했듯이 paired dataset의 경우, X와 Y의 이미지의 좌표값이 있어서 이 관계에 대한 정보가 있습니다. 반면 unpaired dataset은 이 정보가 없으며, 각 이미지간의 대응관계 너무 많기 때문에 사실상 만들어진 이미지가 실제와 pair라고 확정지을 수 없습니다. 즉, mapping의 제약이 적다는 의미이며 이는 결국 mode collapse를 야기할 수 있습니다. 이를 방지하기위해 Cycle consistency를 사용하는 것입니다.

단순 mapping이 아닌 돌아오는 mapping까지 고려하여, 최대한 mapping의 가능성을 줄여주는 것이죠. 그래서 아래의 loss function이 나오게 된 것입니다.

Adversarial Loss

CycleGAN 역시 생성모델이기에 adversarial loss를 사용합니다. 하지만 여기서 알아둬야할 점은 일반적인 GAN처럼 BCE loss를 사용하지 않고 Least Square loss를 사용한다는 점입니다.

논문에서는 이렇게 Least Sqaure loss를 사용했을 때, 좀 더 안정적으로 학습이 가능했고, 높은 퀄리티의 이미지를 생성했다고 합니다. 그리고 이 loss를 사용하면 BCE loss를 사용했을 때 나타나는 mode collapse나 vanishing gradient의 가능성을 줄여준다고 합니다.

Identity Loss

마지막으로 Identity Loss입니다. 논문의 저자들은 일부 작업에 이 loss를 적용했습니다. 일반적으로 모델은 X에서 Y로 가는 mapping을 학습합니다. 아래의 이미지가 대표적인 결과물입니다.

input은 모네의 그림인데요. 그림이 들어오면 마치 사진처럼 굉장히 잘 바꾼다는 것은 확실합니다. 하지만 색상이나 분위기를 완전히 바뀌었습니다. 이처럼 분위기나 색상을 유지하기 위해 Identity loss를 사용한 것입니다.

즉, Identity loss를 사용하는 이유는 Y에서 Y로 가는 mapping에 대한 제약을 걸어주는 것입니다. Y domain의 image가 input으로 들어왔을 때 그 특징을 유지하려면, 모델은 단순히 X를 Y로 만드는 단순 작업(?)을 하는 것이 아니라 들어온 input 이미지를 한 번 더 들여다 봐야한다는 것이죠.

그래서 위와 같은 identity loss를 만들어주는 것입니다.

    if gan_mode == 'lsgan':
        self.loss = nn.MSELoss()
    elif gan_mode == 'vanilla':
        self.loss = nn.BCEWithLogitsLoss()
    elif gan_mode in ['wgangp']:
        self.loss = None
    else:
        raise NotImplementedError('gan mode %s not implemented' % gan_mode)

---

        self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
        self.criterionCycle = torch.nn.L1Loss()
        self.criterionIdt = torch.nn.L1Loss()

한계

지금까지 본 CycleGAN은 굉장히 좋은 성능을 보여줬습니다. 하지만 이 모델 역시 한계는 존재했습니다. 우선 이미지의 모양을 바꾸진 못한다는 점입니다. CycleGAN은 주로 분위기나 색상을 바꾸는 것으로 스타일을 학습하여 다른 이미지를 생성합니다. 그러다보니 피사체의 모양 자체는 바꿀 수가 없었습니다. 그리고 데이터셋의 분포가 불안정하면 이미지를 제대로 생성할 수 없게 됩니다.

위의 예시처럼 사과를 오렌지로 바꾸는 작업은 모양을 바꾸지 못하는 오류입니다. 그리고 사람을 태운 말을 얼룩말로 바꿀 때, 사람까지 얼룩말이 되는 것은 데이터셋에 사람이 말을 탄 데이터가 없기 때문에 나타난 결과입니다.

Reference

CycleGAN, pix2pix 논문 :
https://arxiv.org/abs/1611.07004
https://arxiv.org/abs/1703.10593
coursera 강의 : Apply Generative Adversarial Networks (GANs)
pix2pix 관련 :
https://taeoh-kim.github.io/blog/image2image/
https://greeksharifa.github.io/generative%20model/2019/04/07/Pix2Pix/
https://www.secmem.org/blog/2020/07/19/pix2pix/
CycleGAN 관련 :
https://taeoh-kim.github.io/blog/image2image/
https://comlini8-8.tistory.com/9
https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
https://medium.com/curg/cyclegan-unpaired-%EB%8D%B0%EC%9D%B4%ED%84%B0%EB%A5%BC-%ED%95%99%EC%8A%B5%ED%95%98%EA%B3%A0-%EC%9D%B4%EB%AF%B8%EC%A7%80-%EB%B3%80%ED%99%98%ED%95%98%EA%B8%B0-6fca2f2cddd5