‘Pixel Recurrent Neural Networks‘[1601.06759] 페이퍼에는 세개의 네트워크 모델을 설명하고 있습니다. PixelRNN 으로 명명된 Row LSTM, Diagonal BiLSTM 과 PixelCNN 입니다. 이 중에 PixelCNN 부분을 재현해 보았습니다. PixelCNN 은 PixelRNN 에 비해 비교적 학습 속도가 빠르고 PixelCNN 후속 페이퍼도 계속 나오고 있습니다. 문제를 단순화하기 위해 MNIST 손글씨 숫자 흑백 이지미를 대상으로 하였습니다.

PixelCNN 의 몇가지 주요한 핵심 중 하나는 입력 특성을 콘볼루션할 때 필터의 일부분만 사용한다는 점입니다. 입력 특성맵과 출력 특성맵이 같은 크기이고(텐서플로우에서 ‘SAME’ 패딩) 스트라이드가 1일 때 홀수 크기의 필터는 필터의 중심이 입력 특성맵의 모든 픽셀을 지나가게 됩니다. 이 때 필터의 중심에서 오른쪽 그리고 아래 부분의 가중치를 0 으로 셋팅하면 입력 특성맵에서 필터의 중심의 왼쪽과 위쪽의 픽셀에만 가중치를 적용(콘볼빙)하게 됩니다. 이를 그림으로 나타낸 것이 페이퍼에 잘 나와 있습니다.

왼쪽 그림에서 아래 레이에에서 콘볼빙하여 윗 레이어의 특성 맵을 만들 때 아래 레이어의 3 x 3 필터 중심(붉은색)의 왼쪽과 위의 픽셀에만 가중치가 곱해지는 것을 보이고 있습니다. 이를 마스크 콘볼루션(masked convolution)이라고 부르고 있습니다. 필터를 마스크하기 위해서는 필터의 중심에서 우측과 아래의 가중치를 0 으로 셋팅하면 됩니다.

페이퍼에서는 필터의 종류를 두가지로 나누어 적용하였습니다. A 타입은 필터의 중심도 가중치를 0 으로 셋팅하는 것이고 B 타입은 필터의 중심의 값을 이용하는 경우입니다. A 타입은 입력 이미지에 대해 첫번째 콘볼빙할 때만 사용하고 그 이후의 콘볼루션에서는 모두 B 타입을 사용했습니다. 흑백이 아니고 컬러 이미지일 경우 RGB 컬러가 각각 계산되어야 합니다. 즉 픽셀 레벨이 아니고 채널 레벨로 콘볼루션이 됩니다. 컬러일 경우 A 타입은 현재 컬러 채널은 사용하지 않고 B 타입은 자기 컬러 채널도 콘볼루션에 이용하는 식입니다. 여기서는 흑백 이미지를 사용하므로 컬러 채널에 대해서는 고려하지 않겠습니다.

마스크 콘볼루션을 위해 텐서플로우의 tf.nn.conv2d 함수에 마스크된 weight 를 넘기기 위한 conv2d 함수를 만들어 사용하겠습니다.

def conv2d(input, num_output, kernel_shape, mask_type, scope="conv2d"):

    with tf.variable_scope(scope):

        kernel_h, kernel_w = kernel_shape
        center_h = kernel_h // 2
        center_w = kernel_w // 2

        channel_len = input.get_shape()[-1]
        mask = np.ones((kernel_h, kernel_w, channel_len, num_output), dtype=np.float32)
        mask[center_h, center_w+1:, :, :] = 0.
        mask[center_h+1:, :, :, :] = 0.
        if mask_type == 'A':
            mask[center_h, center_w, :, :] = 0.

        weight = tf.get_variable("weight", [kernel_h, kernel_w, channel_len, num_output], tf.float32, tf.contrib.layers.xavier_initializer())
        weight *= tf.constant(mask, dtype=tf.float32)

        value = tf.nn.conv2d(input, weight, [1, 1, 1, 1], padding='SAME', name='value')
        bias = tf.get_variable("bias", [num_output], tf.float32, tf.zeros_initializer)
        output = tf.nn.bias_add(value, bias, name='output')

        print('[conv2d_%s] %s : %s %s -> %s %s' % (mask_type, scope, input.name, input.get_shape(), output.name, output.get_shape()))

        return output

이 코드에서 핵심은 mask 넘파이(numpy) 배열입니다. 이 배열을 가중치와 같은 사이즈로 만든 다음, 가운데 왼쪽과 위쪽만 1 값을 남기도록 하고 있습니다. 그런 다음 이 배열을 가중치 weight 와 곱하면 가중치 배열의 중앙 왼쪽과 위쪽만 남게 됩니다. 마스크 타입이 A 일 경우에는 정중앙의 마스크 값고 0 으로 셋팅하고 있습니다. 그리고 ‘SAME’ 패딩을 사용해서 입력과 출력 이미지의 크기를 동일하게 맞추고 있습니다. 이는 나중에 학습된 모델로 동일한 크기의 이미지를 생성해 내기 위해서입니다.

페이퍼에서는 PixelCNN 의 필터 사이즈는 첫번째 레이어에서는 7 x 7 과 그 다음 부터는 3 x 3 크기를 사용했습니다. 그리고 총 15개의 레이어와 특성 맵의 수는 128 개를 사용했다고 합니다. 하지만 이와 같은 크기로 학습시키기에는 무리가 있어 레이어는 6 개 정도로 줄이고 특성 맵은 64개를 사용해서 테스트하였습니다.

hidden_layer = [conv2d(X, 64, [7, 7], "A", scope="conv_A")]
for i in range(2):
    hidden_layer.append(conv2d(hidden_layer[-1], 64, [3, 3], "B", scope='conv_B_%d' % i))

for i in range(2):
    hidden_layer.append(tf.nn.relu(conv2d(hidden_layer[-1], 64, [1, 1], "B", scope='relu_B_%d' % i)))

y_logits = conv2d(hidden_layer[-1], 1, [1, 1], "B", scope='y_logits')
y_ = tf.nn.sigmoid(y_logits)

이 코드에서 첫번째 콘볼루션은 7 x 7 필터 크기에 마스크 타입 A 를 사용했습니다. 그 다음 두개 콘볼루션은 3 x 3 에 마스크 타입 B 를 사용했습니다. 그 다음의 두개 콘볼루션은 1 x 1 필터에 렐루 활성화 함수를 거쳐서 마지막 콘볼루션에서는 출력 맵 사이즈를 1 로 만들어 64개의 특성 맵을 하나로 줄였습니다. 흑백 이미지를 사용하므로 0, 1 문제에 맞는 시그모이드 함수를 사용했습니다. 컬러 채널일 경우에는 소프트맥스(softmax) 함수를 사용합니다.

for epoch in range(300):
    # train
    total_train_costs = []
    for i in range(train_step):
        batch = mnist.train.next_batch(100)
        images = binarize(batch[0]).reshape([100, 28, 28, 1])
        _, cost = sess.run([optim, loss], feed_dict={X: images})
        total_train_costs.append(cost)

    # test
    ...

    # generate samples
    samples = np.zeros((100, 28, 28, 1), dtype='float32')
    for i in range(28):
        for j in range(28):
            for k in range(1):
                next_sample = binarize(sess.run(y_, {X: samples}))
                samples[:, i, j, k] = next_sample[:, i, j, k]

    samples = samples.reshape((10, 10, 28, 28))
    samples = samples.transpose(1, 2, 0, 3)
    samples = samples.reshape((28 * 10, 28 * 10))

    toimage(samples, cmin=0.1, cmax=1.0).save('sample/epoch_%d.jpg' % (epoch+1))

학습 과정은 총 300 번의 반복(epoch)를 수행합니다. MNIST 훈련 데이터를 100개씩 미니배치로 학습시키며 최적화 알고리즘은 Adam 옵티마이저를 사용했습니다. 매 반복마다 학습된 모델로 이미지를 생성해서 학습의 진행 정도에 따라 어떤 모델이 만들어지는지 확인해 보았습니다. 이미지를 생성할 때 총 100개의 28 x 28 x 1 사이즈의 빈 배열을 만들어 100개의 이미지를 만들도록 주입합니다. 이미지는 모두 0 으로 채워진 빈 28 x 28 픽셀에서 맨 좌측 상단 부터 한 픽셀씩 채워지도록 합니다. 이렇게 하기 위해서 sess.run 에서 리턴 받은 이미지에서 한 픽셀씩 순서대로 사용하고 나머지 픽셀은 버립니다(약간 비효율적이긴 하지만 텐서플로우의 콘볼루션 기능을 그대로 사용하기 위해서는 어쩔 수 없습니다).

MNIST 데이터는 이미지의 흑백 정도를 0 에서 1 사이의 값으로 표현합니다. 이미지의 컬러가 불연속적인 값으로 표현할 때 학습에도 유리하고 또 더 나은 결과를 준다고 합니다. 즉 흑백 이미지면 0 또는 1 이고 컬러일 경우는 0 ~ 255 까지의 값입니다. 여기에서는 흑백이미지 이므로 binarize 함수에서 랜덤한 균등분포를 만들어 픽셀의 값을 0 또는 1 로 랜덤하게 설정합니다. 100개의 이미지를 10 x 10 격자 형태로 저장하려고 배열의 모양을 바꾸어 280 x 280 하나의 배열로 만들었습니다.

이미지를 저장하기 위해서는 scipy 의 toimage 함수를 사용했습니다. 이 함수를 사용하기 위해서는 파이썬 이미지 라이브러리인 PIL 이 설치되어 있어야 합니다. PIL 은 파이썬 3.x 을 아직 지원하지 않고 있으므로 PIL 의 파이썬 3.x 지원 포크인 Pillow 를 설치하면 됩니다. 아나콘다를 사용할 경우 conda install pillow 로 쉽게 설치할 수 있습니다.

아래는 반복 1, 100, 200, 300 번째에서 생성한 손글씨 숫자 이미지입니다.

전체 코드는 깃허브에서 확인하실 수 있습니다. 컴퓨터 사양에 따라 학습에 시간이 많이 소요될 경우에는 모델을 반복마다 저장해 가면서 학습을 나눠서 진행하는 것이 좋습니다.

참고자료

• carpedm20/pixel-rnn-tensorflow
• igul222/pixel_rnn
• tensorflow/magenta