실전 - CycleGAN 모델 구현하기

CycleGAN은 GAN 영역에서 매우 유명한 모델이라고 할 수 있습니다. 이전에 연구된 pix2pix 같은 두 클래스 간 이미지들의 변환에는 pair 간 학습이 이루어져야 했지만, CycleGAN에서는 unpaired된 상태에서도 학습을 수행하기 위해 cycle consistency를 도입하였습니다. 자세한 설명은 논문과 구글링을 참고하시기 바랍니다. 이 포스트에서는 논문에서의 네트워크 구조로부터 모델을 PyTorch로 구현해 보겠습니다.

전체적인 Generator 구조 파악하고 구현하기

GAN은 Generator와 Discriminator 두 모델로 구성되어 있고, 각 모델에 대해 손실을 최적화하는 방식으로 학습이 진행됩니다. 먼저 Generator에 대한 구조를 논문으로부터 살펴보겠습니다.

논문으로부터 고려해야 할 사항들은 다음과 같습니다.

• 먼저, 사용한 생성자는 입력되는 이미지의 해상도에 따라 residual blocks가 다른 것을 알 수 있습니다. 이를 Generator 생성자의 인자로 지정할 수 있도록 하겠습니다.

• Rk로 Residual block을 설계합니다. 이는 다른 모듈로 정의하겠습니다.

• dk와 uk는 feature map을 줄이거나 늘리는 레이어입니다. 컨볼루션 레이어에 사용되는 종류만 다른 점만 고려하면 되겠습니다.

이들을 고려해서, Generator 클래스로 전체적인 네트워크 구조를 중심으로 구현해 보겠습니다.

from torch import nn

class Generator(nn.Module):
    def __init__(self, num_blocks):
        super().__init__()
        model = []

        # 1, c7s1-64
        model += [
            nn.ReflectionPad2d(3),
            nn.Conv2d(3, 64, 7),
            nn.InstanceNorm2d(64),
        ]

        # 2, dk
        model += [self.__conv_block(64, 128), self.__conv_block(128, 256)]

        # 3, Rk
        model += [ResidualBlock()] * num_blocks

        # 4, uk
        model += [
            self.__conv_block(256, 128, upsample=True),
            self.__conv_block(128, 64, upsample=True),
        ]

        # 5, c7s1-3
        model += [nn.ReflectionPad2d(3), nn.Conv2d(64, 3, 7), nn.Tanh()]

        # 6
        self.model = nn.Sequential(*model)

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

1. Convolution-InstanceNorm-ReLU layer을 리스트로 추가해줍니다. 본문 중에 Reflection padding을 사용했다고 하는데, 이를 nn.ReflectionPad2d으로 미리 늘린 다음에, 컨볼루션 레이어에는 padding은 사용하지 않도록 합니다.

2. 두 개의 dk 레이어로 크기를 줄입니다.

3. 여러 개의 Rk, 즉 Residual block을 추가해줍니다. 여기서 블록 수는 생성자의 매개변수로 사용하도록 합니다.

4. uk를 dK와 같은 메서드로 생성하도록 합니다. 그런데 여기서는 fractional-strided-Convolution을 사용하도록 되어 있으므로, 크기를 늘리는 점을 명시하기 위해 upsample 인자로 전달하겠습니다.

5. 마지막 블록인 c7s1-3입니다. 참고로 제일 마지막 레이어에는 Tanh 함수를 사용하는데, 이는 이미지의 값을 -1에서 1로 유지하기 위해서입니다(참고). 추가적으로 정규화 레이어도 마지막 레이어단에서는 일반적으로 사용되지 않습니다.

6. 리스트에서 nn.Sequential로 순서대로 통과해주게 만듭니다. 여기서 리스트에 *를 사용해서 리스트가 아니라 각 값이 인자로 들어가게 해줍니다.

7. nn.Sequential로 만든 모델을 통과합니다.

비슷한 구조를 가지는 블록 구현하기

전체적인 구조에 대해 구현을 완료하였습니다. uk, dk는 컨볼루션 레이어 외에는 구조가 비슷하고, residual block도 하나를 구현해서 선언해 사용하면 되므로 분리했습니다. 이에 해당하는 __conv_block 메서드와 ResidualBlock 클래스를 구현하겠습니다.

먼저 Generator 내에 있는 __conv_block입니다.

    def __conv_block(self, in_features, out_features, upsample=False):
        if upsample:
            # 8
            conv = nn.ConvTranspose2d(
                in_features, out_features, 3, 2, 1, output_padding=1
            )
        else:
            conv = nn.Conv2d(in_features, out_features, 3, 2, 1)

        # 9
        return nn.Sequential(
            conv,
            nn.InstanceNorm2d(256),
            nn.ReLU(),
        )

8. feature map의 크기를 늘릴 때는 nn.Transpose2d를 사용합니다. 그리고, 크기를 줄일 때와 다르게 padding뿐만 아니라 output_padding도 적용해야 정확하게 2배로 늘려집니다.

9. nn.Sequential로 정의한 conv와 함께 반환해줍니다.

다음은 ResidualBlock입니다. 이름에서 드러나듯 마지막 레이어의 출력과 입력을 더해 준다는 특징이 있습니다.

class ResidualBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_block = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(256, 256, 3),
            nn.InstanceNorm2d(256),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(256, 256, 3),
            nn.InstanceNorm2d(256),
        )

    def forward(self, x):
        return x + self.conv_block(x)

별다른 특징은 없고 forward 함수에서 입력 텐서와 더해준다는 것만 잊지 않으면 됩니다.

Discriminator 구현하기

다음은 구별자입니다. 먼저 구조를 살펴보겠습니다.

논문에서 적혀있다 싶이 구별자는 PatchGAN의 구조로 되어있다고 합니다. 참고로 여기서의 70x70은 결과 텐서의 크기가 아닌 receptive field인 것으로 보면 되겠습니다. Ck가 반복됨을 이용해서 Discriminator 클래스를 구현해 보겠습니다.

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()

        # 10
        self.model = nn.Sequential(
            self.__conv_layer(3, 64, norm=False),
            self.__conv_layer(64, 128),
            self.__conv_layer(128, 256),
            self.__conv_layer(256, 512, stride=1),
            nn.Conv2d(512, 1, 4, 1, 1),
        )

    # 11
    def __conv_layer(self, in_features, out_features, stride=2, norm=True):
        layer = [nn.Conv2d(in_features, out_features, 4, stride, 1)]

        if norm:
            layer.append(nn.InstanceNorm2d(out_features))

        layer.append(nn.LeakyReLU(0.2))

        layer = nn.Sequential(*layer)

        return layer

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

10. 바로 nn.Sequential을 사용했습니다. Ck를 __conv_layer 메서드로 구현하게 됩니다. C64에는 정규화를 사용하지 않기 때문에, norm 인자에 명시해 주고, 논문에 써있지는 않지만 C512는 stride가 1인 점도 인자를 통해 전달하도록 합니다(참고).

11. 매개변수로 받은 입출력 feature map의 갯수, stride, 정규화 유뮤에 따라 모델의 구조를 다르게 해줍니다. 리스트로 먼저 모듈들을 넣어준 다음에 nn.Sequential을 사용하는 방식입니다.

구현된 모델 확인하기

모델을 모두 구현했으면, 잘 돌아가는지 확인할 필요가 있습니다. 보통 테스트하기 위해서, 임의의 텐서를 넣고 결과가 나오는지 먼저 확인하고, 그 다음에 텐서의 shape가 원하는 대로 나오는지 체크하는 방식입니다.

모델을 테스트하는 방법으로는 크게 두 가지가 있습니다. 한 모델 안에 여러 모듈이 존재하는데, 각 작은 모듈들을 구현한 뒤 테스트해서 큰 모듈, 모델을 테스트하는 bottom-top 방법과, 전체 모델을 통과해본 뒤 문제가 생긴 부분에 대해 해당하는 모듈을 분석하는 top-bottom 방법이 있습니다. 모델의 크기가 커질수록 bottom-top 방법을 사용하는 것을 추천합니다.

그러면 Generator와 Discriminator 모델을 동작하는지 확인해보겠습니다.

if __name__ == "__main__":
    x = torch.rand((1, 3, 256, 256))
    generator = Generator(6)
    discriminator = Discriminator()

    print("G(x) shape:", generator(x).shape)
    print("D(x) shape:", discriminator(x).shape)

여기서는 256x256 이미지를 입력으로 넣는다고 가정하겠습니다. 이 때 사용하는 residual block의 수는 6개이므로, Generator 생성자의 인수로 6을 넣어줍니다. 그러면 해당 코드를 실행해보겠습니다.

G(x) shape: torch.Size([1, 3, 256, 256])
D(x) shape: torch.Size([1, 1, 30, 30])

결과적으로 두 모델 모두 shape는 잘 나온 것을 확인할 수 있습니다.

참고
https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
https://github.com/aitorzip/PyTorch-CycleGAN