날짜를 저장하십시오! Google I / O가 5 월 18 일부터 20 일까지 반환됩니다. 지금 등록

Pix2Pix

TensorFlow.org에서 보기 Google Colab에서 실행하기 GitHub에서소스 보기 노트북 다운로드하기

이 노트북은 Conditional Adversarial Networks를 사용한 이미지 간 변환에서 설명한 대로 조건부 GAN을 사용한 이미지간 변환을 보여줍니다. 이 기법을 사용하여 흑백 사진을 채색하고, 구글 지도를 구글 어스로 변환하는 등의 작업을 수행할 수 있습니다. 여기서는 건물 정면을 실제 건물로 변환합니다.

예제에서는 프라하 체코 공과대학기계 인식 센터에서 제공하는 CMP Facade Database를 사용할 것입니다. 예제를 간략하게 유지하기 위해 위의 논문 작성자가 만든 이 데이터세트의 전처리된 사본을 사용합니다.

각 epoch는 단일 V100 GPU에서 약 15초가 걸립니다.

다음은 200개 epoch 동안 모델을 훈련한 후에 생성된 출력입니다.

샘플 출력 _1샘플 출력 _2

TensorFlow 및 기타 라이브러리 가져오기

import tensorflow as tf

import os
import time

from matplotlib import pyplot as plt
from IPython import display
pip install -q -U tensorboard

데이터세트 로드하기

이 데이터세트 및 이와 유사한 데이터세트는 여기에서 다운로드할 수 있습니다. 논문에서 언급했듯이 훈련 데이터세트에 무작위 지터링 및 미러링을 적용합니다.

  • 무작위 지터링에서 이미지의 크기는 286 x 286로 조정되고 256 x 256로 무작위로 잘립니다.
  • 무작위 미러링에서는 이미지가 좌우로 무작위로 뒤집힙니다.
_URL = 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz'

path_to_zip = tf.keras.utils.get_file('facades.tar.gz',
                                      origin=_URL,
                                      extract=True)

PATH = os.path.join(os.path.dirname(path_to_zip), 'facades/')
Downloading data from https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz
30171136/30168306 [==============================] - 2s 0us/step
BUFFER_SIZE = 400
BATCH_SIZE = 1
IMG_WIDTH = 256
IMG_HEIGHT = 256
def load(image_file):
  image = tf.io.read_file(image_file)
  image = tf.image.decode_jpeg(image)

  w = tf.shape(image)[1]

  w = w // 2
  real_image = image[:, :w, :]
  input_image = image[:, w:, :]

  input_image = tf.cast(input_image, tf.float32)
  real_image = tf.cast(real_image, tf.float32)

  return input_image, real_image
inp, re = load(PATH+'train/100.jpg')
# casting to int for matplotlib to show the image
plt.figure()
plt.imshow(inp/255.0)
plt.figure()
plt.imshow(re/255.0)
<matplotlib.image.AxesImage at 0x7fa6cbb7b7f0>

png

png

def resize(input_image, real_image, height, width):
  input_image = tf.image.resize(input_image, [height, width],
                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  real_image = tf.image.resize(real_image, [height, width],
                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)

  return input_image, real_image
def random_crop(input_image, real_image):
  stacked_image = tf.stack([input_image, real_image], axis=0)
  cropped_image = tf.image.random_crop(
      stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])

  return cropped_image[0], cropped_image[1]
# normalizing the images to [-1, 1]

def normalize(input_image, real_image):
  input_image = (input_image / 127.5) - 1
  real_image = (real_image / 127.5) - 1

  return input_image, real_image
@tf.function()
def random_jitter(input_image, real_image):
  # resizing to 286 x 286 x 3
  input_image, real_image = resize(input_image, real_image, 286, 286)

  # randomly cropping to 256 x 256 x 3
  input_image, real_image = random_crop(input_image, real_image)

  if tf.random.uniform(()) > 0.5:
    # random mirroring
    input_image = tf.image.flip_left_right(input_image)
    real_image = tf.image.flip_left_right(real_image)

  return input_image, real_image

무작위 지터링을 거치는 아래 이미지에서 볼 수 있듯이 논문의 설명에 따라 무작위 지터링은 다음 작업 수행을 의미합니다.

  1. 더 큰 높이와 너비로 이미지의 크기를 조정합니다.
  2. 대상 크기로 무작위 자릅니다.
  3. 이미지를 가로로 무작위 뒤집습니다.
plt.figure(figsize=(6, 6))
for i in range(4):
  rj_inp, rj_re = random_jitter(inp, re)
  plt.subplot(2, 2, i+1)
  plt.imshow(rj_inp/255.0)
  plt.axis('off')
plt.show()

png

def load_image_train(image_file):
  input_image, real_image = load(image_file)
  input_image, real_image = random_jitter(input_image, real_image)
  input_image, real_image = normalize(input_image, real_image)

  return input_image, real_image
def load_image_test(image_file):
  input_image, real_image = load(image_file)
  input_image, real_image = resize(input_image, real_image,
                                   IMG_HEIGHT, IMG_WIDTH)
  input_image, real_image = normalize(input_image, real_image)

  return input_image, real_image

입력 파이프라인

train_dataset = tf.data.Dataset.list_files(PATH+'train/*.jpg')
train_dataset = train_dataset.map(load_image_train,
                                  num_parallel_calls=tf.data.experimental.AUTOTUNE)
train_dataset = train_dataset.shuffle(BUFFER_SIZE)
train_dataset = train_dataset.batch(BATCH_SIZE)
test_dataset = tf.data.Dataset.list_files(PATH+'test/*.jpg')
test_dataset = test_dataset.map(load_image_test)
test_dataset = test_dataset.batch(BATCH_SIZE)

생성기 빌드하기

  • 생성기 아키텍처는 수정된 U-Net입니다.
  • 인코더의 각 블록은 (Conv -> Batchnorm -> Leaky ReLU)입니다.
  • 디코더의 각 블록은 (Transposed Conv -> Batchnorm -> Dropout(처음 3개의 블록에 적용됨)-> ReLU)입니다.
  • 인코더와 디코더 사이에는 스킵(skip) 연결이 있습니다(U-Net에서와 같이).
OUTPUT_CHANNELS = 3
def downsample(filters, size, apply_batchnorm=True):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
      tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
                             kernel_initializer=initializer, use_bias=False))

  if apply_batchnorm:
    result.add(tf.keras.layers.BatchNormalization())

  result.add(tf.keras.layers.LeakyReLU())

  return result
down_model = downsample(3, 4)
down_result = down_model(tf.expand_dims(inp, 0))
print (down_result.shape)
(1, 128, 128, 3)
def upsample(filters, size, apply_dropout=False):
  initializer = tf.random_normal_initializer(0., 0.02)

  result = tf.keras.Sequential()
  result.add(
    tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                    padding='same',
                                    kernel_initializer=initializer,
                                    use_bias=False))

  result.add(tf.keras.layers.BatchNormalization())

  if apply_dropout:
      result.add(tf.keras.layers.Dropout(0.5))

  result.add(tf.keras.layers.ReLU())

  return result
up_model = upsample(3, 4)
up_result = up_model(down_result)
print (up_result.shape)
(1, 256, 256, 3)
def Generator():
  inputs = tf.keras.layers.Input(shape=[256,256,3])

  down_stack = [
    downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64)
    downsample(128, 4), # (bs, 64, 64, 128)
    downsample(256, 4), # (bs, 32, 32, 256)
    downsample(512, 4), # (bs, 16, 16, 512)
    downsample(512, 4), # (bs, 8, 8, 512)
    downsample(512, 4), # (bs, 4, 4, 512)
    downsample(512, 4), # (bs, 2, 2, 512)
    downsample(512, 4), # (bs, 1, 1, 512)
  ]

  up_stack = [
    upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)
    upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)
    upsample(512, 4), # (bs, 16, 16, 1024)
    upsample(256, 4), # (bs, 32, 32, 512)
    upsample(128, 4), # (bs, 64, 64, 256)
    upsample(64, 4), # (bs, 128, 128, 128)
  ]

  initializer = tf.random_normal_initializer(0., 0.02)
  last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,
                                         strides=2,
                                         padding='same',
                                         kernel_initializer=initializer,
                                         activation='tanh') # (bs, 256, 256, 3)

  x = inputs

  # Downsampling through the model
  skips = []
  for down in down_stack:
    x = down(x)
    skips.append(x)

  skips = reversed(skips[:-1])

  # Upsampling and establishing the skip connections
  for up, skip in zip(up_stack, skips):
    x = up(x)
    x = tf.keras.layers.Concatenate()([x, skip])

  x = last(x)

  return tf.keras.Model(inputs=inputs, outputs=x)
generator = Generator()
tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64)

png

gen_output = generator(inp[tf.newaxis,...], training=False)
plt.imshow(gen_output[0,...])
Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).
<matplotlib.image.AxesImage at 0x7fa65c3a9828>

png

  • 생성기 손실
    • 이것은 생성된 이미지 및 1의 배열의 시그모이드 교차 엔트로피 손실입니다.
    • 논문에는 생성 이미지와 대상 이미지 사이의 MAE(평균 절대 오차)인 L1 손실도 포함됩니다.
    • 이를 통해 생성된 이미지가 대상 이미지와 구조적으로 유사해질 수 있습니다.
    • 총 생성기 손실을 계산하는 공식 = gan_loss + LAMBDA * l1_loss. 여기서 LAMBDA = 100. 이 값은 논문 작성자가 결정했습니다.

생성기의 훈련 절차는 다음과 같습니다.

LAMBDA = 100
def generator_loss(disc_generated_output, gen_output, target):
  gan_loss = loss_object(tf.ones_like(disc_generated_output), disc_generated_output)

  # mean absolute error
  l1_loss = tf.reduce_mean(tf.abs(target - gen_output))

  total_gen_loss = gan_loss + (LAMBDA * l1_loss)

  return total_gen_loss, gan_loss, l1_loss

발전기 업데이트 이미지

판별자 빌드하기

  • 판별자는 PatchGAN입니다.
  • 판별자의 각 블록은 (Conv -> BatchNorm -> Leaky ReLU)입니다.
  • 마지막 레이어 이후의 출력 형상은 (batch_size, 30, 30, 1)입니다.
  • 출력의 각 30x30 패치는 입력 이미지의 70x70 부분을 분류합니다(이러한 아키텍처를 PatchGAN이라고 함).
  • 판별자는 2개의 입력을 받습니다.
    • 진짜로 분류해야 하는 입력 이미지 및 대상 이미지
    • 가짜로 분류해야 하는 입력 이미지 및 생성된 이미지(생성기의 출력)
    • 코드에서 이 두 입력을 함께 연결합니다(tf.concat([inp, tar], axis=-1)).
def Discriminator():
  initializer = tf.random_normal_initializer(0., 0.02)

  inp = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')
  tar = tf.keras.layers.Input(shape=[256, 256, 3], name='target_image')

  x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2)

  down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64)
  down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)
  down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)

  zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
  conv = tf.keras.layers.Conv2D(512, 4, strides=1,
                                kernel_initializer=initializer,
                                use_bias=False)(zero_pad1) # (bs, 31, 31, 512)

  batchnorm1 = tf.keras.layers.BatchNormalization()(conv)

  leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)

  zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)

  last = tf.keras.layers.Conv2D(1, 4, strides=1,
                                kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1)

  return tf.keras.Model(inputs=[inp, tar], outputs=last)
discriminator = Discriminator()
tf.keras.utils.plot_model(discriminator, show_shapes=True, dpi=64)

png

disc_out = discriminator([inp[tf.newaxis,...], gen_output], training=False)
plt.imshow(disc_out[0,...,-1], vmin=-20, vmax=20, cmap='RdBu_r')
plt.colorbar()
<matplotlib.colorbar.Colorbar at 0x7fa65c2862e8>

png

판별자 손실

  • 판별자 손실 함수는 진짜 이미지, 생성된 이미지의 두 입력을 받습니다.
  • real_loss는 진짜 이미지1의 배열(진짜 이미지이기 때문)의 시그모이드 교차 엔트로피 손실입니다.
  • generated_loss는 생성된 이미지0의 배열(가짜 이미지이기 때문)의 시그모이드 교차 엔트로피 손실입니다.
  • 그리고 total_loss는 real_loss와 generated_loss의 합계입니다.
loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)
def discriminator_loss(disc_real_output, disc_generated_output):
  real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output)

  generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output)

  total_disc_loss = real_loss + generated_loss

  return total_disc_loss

판별자의 훈련 절차는 다음과 같습니다.

아키텍처 및 하이퍼 매개변수에 대한 자세한 내용은 논문을 참조하세요.

판별 기 업데이트 이미지

옵티마이저 및 체크포인트-세이버 정의하기

generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

이미지 생성하기

훈련 중에 일부 이미지를 플롯하는 함수를 작성합니다.

  • 테스트 데이터세트의 이미지를 생성기로 전달합니다.
  • 그러면 생성기가 입력 이미지를 출력으로 변환합니다.
  • 마지막 단계로 예상을 플롯합니다. 짜잔!

참고: 여기서 training=True는 테스트 데이터세트에서 모델을 실행하는 동안 배치 통계를 얻으려고 하므로 의도적인 것입니다. training = False를 사용하면 훈련 데이터세트에서 학습한 누적 통계(이것을 원하지는 않음)를 얻게 됩니다.

def generate_images(model, test_input, tar):
  prediction = model(test_input, training=True)
  plt.figure(figsize=(15,15))

  display_list = [test_input[0], tar[0], prediction[0]]
  title = ['Input Image', 'Ground Truth', 'Predicted Image']

  for i in range(3):
    plt.subplot(1, 3, i+1)
    plt.title(title[i])
    # getting the pixel values between [0, 1] to plot it.
    plt.imshow(display_list[i] * 0.5 + 0.5)
    plt.axis('off')
  plt.show()
for example_input, example_target in test_dataset.take(1):
  generate_images(generator, example_input, example_target)

png

훈련하기

  • 각 예에서 입력은 출력을 생성합니다.
  • 판별자는 input_image 및 생성된 이미지를 첫 번째 입력으로 받습니다. 두 번째 입력은 input_image와 target_image입니다.
  • 다음으로 생성기와 판별자 손실을 계산합니다.
  • 그런 다음 생성기와 판별자 변수(입력) 모두에 대해 손실의 그래디언트를 계산하고 이를 옵티마이저에 적용합니다.
  • 그런 다음 손실을 TensorBoard에 기록합니다.
EPOCHS = 150
import datetime
log_dir="logs/"

summary_writer = tf.summary.create_file_writer(
  log_dir + "fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
@tf.function
def train_step(input_image, target, epoch):
  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
    gen_output = generator(input_image, training=True)

    disc_real_output = discriminator([input_image, target], training=True)
    disc_generated_output = discriminator([input_image, gen_output], training=True)

    gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target)
    disc_loss = discriminator_loss(disc_real_output, disc_generated_output)

  generator_gradients = gen_tape.gradient(gen_total_loss,
                                          generator.trainable_variables)
  discriminator_gradients = disc_tape.gradient(disc_loss,
                                               discriminator.trainable_variables)

  generator_optimizer.apply_gradients(zip(generator_gradients,
                                          generator.trainable_variables))
  discriminator_optimizer.apply_gradients(zip(discriminator_gradients,
                                              discriminator.trainable_variables))

  with summary_writer.as_default():
    tf.summary.scalar('gen_total_loss', gen_total_loss, step=epoch)
    tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch)
    tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch)
    tf.summary.scalar('disc_loss', disc_loss, step=epoch)

실제 훈련 루프:

  • epoch 수에 대해 반복합니다.
  • 각 epoch에서 표시 내용을 지우고 generate_images를 실행하여 진행 상황을 표시합니다.
  • 각 epoch마다 훈련 데이터세트에 대해 반복하여 각 예에 대해 '.'를 인쇄합니다.
  • 20개 epoch마다 체크포인트를 저장합니다.
def fit(train_ds, epochs, test_ds):
  for epoch in range(epochs):
    start = time.time()

    display.clear_output(wait=True)

    for example_input, example_target in test_ds.take(1):
      generate_images(generator, example_input, example_target)
    print("Epoch: ", epoch)

    # Train
    for n, (input_image, target) in train_ds.enumerate():
      print('.', end='')
      if (n+1) % 100 == 0:
        print()
      train_step(input_image, target, epoch)
    print()

    # saving (checkpoint) the model every 20 epochs
    if (epoch + 1) % 20 == 0:
      checkpoint.save(file_prefix = checkpoint_prefix)

    print ('Time taken for epoch {} is {} sec\n'.format(epoch + 1,
                                                        time.time()-start))
  checkpoint.save(file_prefix = checkpoint_prefix)

이 훈련 루프는 TensorBoard에서 쉽게 볼 수 있는 로그를 저장하여 훈련 진행 상황을 모니터링합니다. 로컬에서 작업하면 별도의 TensorBoard 프로세스가 시작됩니다. 노트북에서 TensorBoard로 모니터링하려는 경우 훈련을 시작하기 전에 뷰어를 시작하는 것이 편합니다.

뷰어를 시작하려면 다음을 코드 셀에 붙여넣습니다.

%load_ext tensorboard
%tensorboard --logdir {log_dir}

이제 훈련 루프를 실행합니다.

fit(train_dataset, EPOCHS, test_dataset)

png

Epoch:  149
....................................................................................................
....................................................................................................
....................................................................................................
....................................................................................................

Time taken for epoch 150 is 16.547991037368774 sec

TensorBoard 결과를 공개적으로 공유하려면 다음을 코드 셀에 복사하여 TensorBoard.dev에 로그를 업로드할 수 있습니다.

참고: 이를 위해 구글 계정이 필요합니다.

tensorboard dev upload --logdir  {log_dir}

주의: 이 명령은 종료되지 않습니다. 장기 실행되는 실험 결과를 지속적으로 업로드하도록 설계되었습니다. 데이터가 업로드되면 노트북 도구에서 "실행 중단" 옵션을 사용하여 실행을 중단시켜야 합니다.

TensorBoard.dev에서 이 노트북의 이전 실행 결과를 볼 수 있습니다.

TensorBoard.dev는 ML 실험을 호스팅 및 추적하고 모든 사람과 공유하기 위한 관리 환경입니다.

<iframe>을 사용하여 인라인으로 포함할 수도 있습니다.

display.IFrame(
    src="https://tensorboard.dev/experiment/lZ0C6FONROaUMfjYkVyJqw",
    width="100%",
    height="1000px")

GAN의 로그 해석은 단순한 분류 또는 회귀 모델보다 까다롭습니다. 살펴 보아야 할 내용::

  • 어떤 모델에도 "won"이 없는 것을 확인합니다. gen_gan_loss 또는 disc_loss가 매우 낮아지면 이 모델이 다른 모델을 지배하고 있고 결합된 모델을 성공적으로 훈련하지 못하고 있음을 나타냅니다.
  • log(2) = 0.69 값은 판별자가 두 옵션에 대해 평균적으로 동일하게 확신하지 못함을 의미하는 2의 perplexity를 나타내기 때문에 이러한 손실에 대한 좋은 기준점입니다.
  • disc_loss의 경우 0.69 미만의 값은 진짜 + 생성된 이미지의 결합된 세트에서 판별자가 무작위보다 성능이 좋음을 나타냅니다.
  • gen_gan_loss의 경우 0.69 미만의 값은 생성기 i가 무작위보다 판별자를 더 잘 속인다는 것을 의미합니다.
  • 훈련이 진행됨에 따라 gen_l1_loss가 줄어듭니다.

최신 체크포인트를 복원하고 테스트

ls {checkpoint_dir}
checkpoint          ckpt-5.data-00000-of-00001
ckpt-1.data-00000-of-00001  ckpt-5.index
ckpt-1.index            ckpt-6.data-00000-of-00001
ckpt-2.data-00000-of-00001  ckpt-6.index
ckpt-2.index            ckpt-7.data-00000-of-00001
ckpt-3.data-00000-of-00001  ckpt-7.index
ckpt-3.index            ckpt-8.data-00000-of-00001
ckpt-4.data-00000-of-00001  ckpt-8.index
ckpt-4.index
# restoring the latest checkpoint in checkpoint_dir
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7fa12e5cf3c8>

테스트 데이터세트를 사용하여 생성하기

# Run the trained model on a few examples from the test dataset
for inp, tar in test_dataset.take(5):
  generate_images(generator, inp, tar)

png

png

png

png

png