[ML Practice] pix2pix(1)

Original article was published by 김현우 on Deep Learning on Medium


휴먼스케이프 Software engineer Covy입니다.

본 포스트에서는 이전 포스트에서 리뷰를 진행했던 pix2pix에서 제시한 layer과 model을 PyTorch를 이용해서 직접 구현해보는 시간을 가지려고 합니다. 이전 포스트가 궁금하신 분들은 이곳을 참고하시면 좋습니다.

다음 포스트에서는 이어서 pix2pix의 training 부분을 구현해보는 시간을 가지겠습니다.

Objective

본 포스트에서는 PyTorch를 이용해 pix2pix 논문에 적혀있는 network의 구조를 구현하는 방법에 대해서 서술합니다. 이는 앞선 포스트에서 소개드렸던 논문의 부록에 적혀있는 구조를 바탕으로 구현한 것입니다.

Network Architecture

논문의 부록에서는 네트워크 아키텍쳐를 구현하기에 앞서 아키텍쳐를 이루는 layer의 집합에 대한 일반적인 형태를소개하고 있습니다.

그 두 가지 형태가 Ck와 CDk입니다.

먼저 Ck의 경우, k만큼의 output channel수를 가진 Convolution-BatchNormalization-ReLU의 형태로 이루어진 layer의 집합체입니다.

다음으로 CDk의 경우, 마찬가지로 k만큼의 output channel수를 가지지만, Ck와는 다르게 Convolution-BatchNormalization-ReLU-Dropout의 형태로 이루어진 layer의 집합체입니다.

이를 바탕으로 논문에서 사용한 아키텍쳐는 다음과 같습니다.

1. Generator(Encoder)
C64-C128-C256-C512-C512-C512-C512-C512

2. Generator(Decoder)
CD512-CD1024-CD1024-C1024-C1024-C512-C256-C128
(이 경우 U-Net의 형태로 i번째 layer 집합체와 8-i번째 layer 집합체간의 skip connection이 존재합니다.)

3. Discriminator
C64-C128-C256-C512

위 세 가지 형태의 아키텍처를 각각 구현하기 위해서 layer.py에 layer집합체인 Ck와 CDk를, model.py에 앞서 정의한 Ck와 CDk를 이용한 generator와 discriminator를 구현하도록 하겠습니다.

Network Implementation

1. CBR2D 구현
먼저 Convolution-BatchNormalization-ReLU를 구현한 layer 집합체인 CBR2D를 구현한 것은 다음과 같다.

우선 parameter로 받는 것들은 in_channels, out_channels, kernel_size, stride, padding, bias인데, 이들 모두 CBR2D의 convolution layer를 정의하기 위한 parameter이다.

__init__에서 nn.Module의 __init__을 불러준 뒤에, layer 집합체에 속할 것들을 layers에 순서대로 모아줍니다. 위 코드에서는 nn.Conv2d, nn.BatchNorm2d 혹은 nn.InstanceNorm2d, 그리고 nn.ReLU 혹은 nn.LeakyReLU 가 layers에 순차적으로 쌓이게 됩니다. 이후, nn.Sequential를 이용하여 순서대로 실행을 담당하는 함수를 클래스 변수인 cbr에 저장합니다.

forward()에서는 앞서 저장한 layer의 순차적 실행(forward propagation을 의미합니다.)을 주어진 input x에 대해서 진행합니다. 이 forward function의 정의는 pyTorch module 설계에서 module name만으로 network의 forward propagation output을 얻어낼 수 있게 합니다. (ex. x가 input일 경우, CBR2d(x) 만으로 forward propagation의 결과를 얻을 수 있습니다.)

2. DECBR2D 구현
다음으로 ConvolutionTranspose-BatchNormalization-ReLU를 구현한 layer 집합체인 DECBR2D를 구현한 것은 다음과 같다.

앞서 설명한 CBR2D와 크게 다를 바가 없는 형태지만, 이 형태를 구현한 이유는 U-Net 형태의 generator에서 decoder를 구현하기 위해서 입니다. CBR2D는 encoder를 위한 layer 집합체였다면, 지금 구현한 DECBR2D는 decoder를 위한 layer의 집합체입니다.

차이를 보시면 알 수 있듯이, convolution layer 대신에 convolutionTranspose layer가 들어가서 전반적으로 줄어든 dimension을 처음 input으로 들어온 dimension만큼으로 늘려주는 작업을 하는 것입니다.

3. Generator 구현
다음으로 앞서 구현한 CBR2D와 DECBR2D를 이용하여 Pix2Pix의 generator를 구현해보겠습니다.

기본적인 구현방법은 앞서 설명했던 것들과 크게 다르지는 않습니다. 이미지의 input_channel 수가 3이기 때문에 default를 3으로 진행했고, pix2pix에서 encoder의 첫 형태가 C64이고 2배씩 늘어나는 형태이기 때문에 default nker를 64로 설정하고 2배씩 늘려주는 형태로 layer들을 구현했습니다.

논문에 적혀있는대로, 4×4 size의 kernel과 stride 2를사용했고, 이러한 조건에서 size를 1/2씩으로 줄이기 위해 padding으로 1을 사용했습니다. 첫 번째와 마지막 layer에서만 normlaization을 제외하였고, 마지막 layer에서는 추가적으로 relu도 제외했습니다.

또한 앞서 decoder에서 CDk형태로 dropout이 존재하는 layer 집합체들이 존재하기 때문에 decoder의 3번째 layer 집합체까지 dropout(0.5)를 정의해주었습니다.

forward()에서는 앞서 정의한 encoder와 decoder를 순차적으로 진행하는데, 여기서 앞서 설명한 skip connection을 구현하기 위해서 torch.cat을 이용하여 i 번째와 8-i 번째의 layer 집합체의 결과를 concatenate해 주면서 propagation을 진행합니다. 더불어 decoder의 마지막 layer가 끝난 뒤에 논문에서 hypertangent를 추가한 것에 따라서 tanh를 사용해주었습니다.

4. Discriminator 구현
Discriminator는 앞선 generator보다 간단한 형태입니다.

마찬가지로, 논문에 적혀있는 내용에 따라서 C64-C128-C256-C512의 구조를 구현한 것입니다.

특징적으로 discriminator에서 다른 점은 마지막 layer의 끝에 sigmoid가 등장한다고 논문에 적혀있으며, activation functio이 leakyReLU 0.2로 주어진다는 점이었습니다.

Conclusion

이것으로 간단하게 pix2pix 논문에서 구현한 model을 이들을 구현하는 general한 layer 집합체인 CBR2D와 DECBR2D를 이용하여 구현해보았습니다. 다음 시간에는 구현한 model을 통해서 학습을 진행하는 방법에 대해서 진행하도록 하겠습니다.

Get to know us better!
Join our official channels below.

Telegram(EN) : t.me/Humanscape
KakaoTalk(KR) : open.kakao.com/o/gqbUQEM
Website : humanscape.io
Medium : medium.com/humanscape-ico
Facebook : www.facebook.com/humanscape
Twitter : twitter.com/Humanscape_io
Reddit : https://www.reddit.com/r/Humanscape_official
Bitcointalk announcement : https://bit.ly/2rVsP4T
Email : support@humanscape.io