`Model.train_step`을 오버라이딩 하는 WGAN-GP

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
import tensorflow as tf
from keras import layers

IMG_SHAPE = (28, 28, 1)
BATCH_SIZE = 512

# 노이즈 벡터의 크기
noise_dim = 128

fashion_mnist = keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
print(f"Number of examples: {len(train_images)}")
print(f"Shape of the images in the dataset: {train_images.shape[1:]}")

# 각 샘플을 (28, 28, 1)로 reshape하고,
# [-1, 1] 범위로 픽셀 값을 정규화합니다.
train_images = train_images.reshape(train_images.shape[0], *IMG_SHAPE).astype("float32")
train_images = (train_images - 127.5) / 127.5

def conv_block(
    x,
    filters,
    activation,
    kernel_size=(3, 3),
    strides=(1, 1),
    padding="same",
    use_bias=True,
    use_bn=False,
    use_dropout=False,
    drop_value=0.5,
):
    x = layers.Conv2D(
        filters, kernel_size, strides=strides, padding=padding, use_bias=use_bias
    )(x)
    if use_bn:
        x = layers.BatchNormalization()(x)
    x = activation(x)
    if use_dropout:
        x = layers.Dropout(drop_value)(x)
    return x


def get_discriminator_model():
    img_input = layers.Input(shape=IMG_SHAPE)
    # 입력 이미지 크기를 (32, 32, 1)로 만들기 위해, 입력을 0으로 채웁니다. (제로 패딩)
    x = layers.ZeroPadding2D((2, 2))(img_input)
    x = conv_block(
        x,
        64,
        kernel_size=(5, 5),
        strides=(2, 2),
        use_bn=False,
        use_bias=True,
        activation=layers.LeakyReLU(0.2),
        use_dropout=False,
        drop_value=0.3,
    )
    x = conv_block(
        x,
        128,
        kernel_size=(5, 5),
        strides=(2, 2),
        use_bn=False,
        activation=layers.LeakyReLU(0.2),
        use_bias=True,
        use_dropout=True,
        drop_value=0.3,
    )
    x = conv_block(
        x,
        256,
        kernel_size=(5, 5),
        strides=(2, 2),
        use_bn=False,
        activation=layers.LeakyReLU(0.2),
        use_bias=True,
        use_dropout=True,
        drop_value=0.3,
    )
    x = conv_block(
        x,
        512,
        kernel_size=(5, 5),
        strides=(2, 2),
        use_bn=False,
        activation=layers.LeakyReLU(0.2),
        use_bias=True,
        use_dropout=False,
        drop_value=0.3,
    )

    x = layers.Flatten()(x)
    x = layers.Dropout(0.2)(x)
    x = layers.Dense(1)(x)

    d_model = keras.models.Model(img_input, x, name="discriminator")
    return d_model


d_model = get_discriminator_model()
d_model.summary()

def upsample_block(
    x,
    filters,
    activation,
    kernel_size=(3, 3),
    strides=(1, 1),
    up_size=(2, 2),
    padding="same",
    use_bn=False,
    use_bias=True,
    use_dropout=False,
    drop_value=0.3,
):
    x = layers.UpSampling2D(up_size)(x)
    x = layers.Conv2D(
        filters, kernel_size, strides=strides, padding=padding, use_bias=use_bias
    )(x)

    if use_bn:
        x = layers.BatchNormalization()(x)

    if activation:
        x = activation(x)
    if use_dropout:
        x = layers.Dropout(drop_value)(x)
    return x


def get_generator_model():
    noise = layers.Input(shape=(noise_dim,))
    x = layers.Dense(4 * 4 * 256, use_bias=False)(noise)
    x = layers.BatchNormalization()(x)
    x = layers.LeakyReLU(0.2)(x)

    x = layers.Reshape((4, 4, 256))(x)
    x = upsample_block(
        x,
        128,
        layers.LeakyReLU(0.2),
        strides=(1, 1),
        use_bias=False,
        use_bn=True,
        padding="same",
        use_dropout=False,
    )
    x = upsample_block(
        x,
        64,
        layers.LeakyReLU(0.2),
        strides=(1, 1),
        use_bias=False,
        use_bn=True,
        padding="same",
        use_dropout=False,
    )
    x = upsample_block(
        x, 1, layers.Activation("tanh"), strides=(1, 1), use_bias=False, use_bn=True
    )
    # 이 시점에서 우리는 입력과 같은 모양(32, 32, 1)을 가진 출력을 얻습니다.
    # 우리는 그것을 (28, 28, 1)로 만들기 위해, Cropping2D 레이어를 사용할 것입니다.
    x = layers.Cropping2D((2, 2))(x)

    g_model = keras.models.Model(noise, x, name="generator")
    return g_model


g_model = get_generator_model()
g_model.summary()

class WGAN(keras.Model):
    def __init__(
        self,
        discriminator,
        generator,
        latent_dim,
        discriminator_extra_steps=3,
        gp_weight=10.0,
    ):
        super().__init__()
        self.discriminator = discriminator
        self.generator = generator
        self.latent_dim = latent_dim
        self.d_steps = discriminator_extra_steps
        self.gp_weight = gp_weight

    def compile(self, d_optimizer, g_optimizer, d_loss_fn, g_loss_fn):
        super().compile()
        self.d_optimizer = d_optimizer
        self.g_optimizer = g_optimizer
        self.d_loss_fn = d_loss_fn
        self.g_loss_fn = g_loss_fn

    def gradient_penalty(self, batch_size, real_images, fake_images):
        """그래디언트 페널티(gradient penalty, GP)를 계산합니다.

        이 손실은 보간된 이미지에서 계산되어, 판별자 손실에 추가됩니다.
        """
        # 보간된 이미지를 얻습니다.
        alpha = tf.random.uniform([batch_size, 1, 1, 1], 0.0, 1.0)
        diff = fake_images - real_images
        interpolated = real_images + alpha * diff

        with tf.GradientTape() as gp_tape:
            gp_tape.watch(interpolated)
            # 1. 이 보간된 이미지에 대한 판별자 출력을 가져옵니다.
            pred = self.discriminator(interpolated, training=True)

        # 2. 이 보간된 이미지에 대한 그래디언트를 계산합니다.
        grads = gp_tape.gradient(pred, [interpolated])[0]
        # 3. 그래디언트의 norm을 계산합니다.
        norm = tf.sqrt(tf.reduce_sum(tf.square(grads), axis=[1, 2, 3]))
        gp = tf.reduce_mean((norm - 1.0) ** 2)
        return gp

    def train_step(self, real_images):
        if isinstance(real_images, tuple):
            real_images = real_images[0]

        # 배치 크기를 가져옵니다.
        batch_size = tf.shape(real_images)[0]

        # 각 배치에 대해 원본 논문에 나와 있는 대로 다음 단계를 수행합니다.
        # 1. 생성자를 트레이닝하고, 생성자 손실을 구합니다.
        # 2. 판별자를 트레이닝하고, 판별자 손실을 구합니다.
        # 3. 그래디언트 페널티(GP)를 계산합니다.
        # 4. 이 그래디언트 페널티를 상수 가중치 인자(factor)와 곱합니다.
        # 5. 판별자 손실에 그래디언트 페널티를 추가합니다.
        # 6. 생성자 및 판별자 손실을 손실 딕셔너리로 반환합니다.

        # 먼저 판별자를 트레이닝합니다.
        # 원본 논문에서는 생성자의 한 단계보다 판별자를 `x` 단계 더(일반적으로 5) 트레이닝하는 것을 권장합니다.
        # 여기서는 트레이닝 시간을 줄이기 위해, 5단계가 아닌 3단계를 더 트레이닝합니다.
        for i in range(self.d_steps):
            # 잠재 벡터를 얻습니다.
            random_latent_vectors = tf.random.normal(
                shape=(batch_size, self.latent_dim)
            )
            with tf.GradientTape() as tape:
                # 잠재 벡터에서 가짜 이미지 생성
                fake_images = self.generator(random_latent_vectors, training=True)
                # 가짜 이미지에 대한 로짓을 얻습니다
                fake_logits = self.discriminator(fake_images, training=True)
                # 실제 이미지에 대한 로짓을 얻습니다
                real_logits = self.discriminator(real_images, training=True)

                # 가짜 및 실제 이미지 로짓을 사용하여, 판별자 손실을 계산합니다.
                d_cost = self.d_loss_fn(real_img=real_logits, fake_img=fake_logits)
                # 그래디언트 페널티를 계산합니다
                gp = self.gradient_penalty(batch_size, real_images, fake_images)
                # 원래 판별자 손실에 그래디언트 페널티를 추가합니다.
                d_loss = d_cost + gp * self.gp_weight

            # 판별자 손실에 대한 그래디언트를 얻습니다.
            d_gradient = tape.gradient(d_loss, self.discriminator.trainable_variables)
            # 판별자 옵티마이저를 사용하여, 판별자의 가중치를 업데이트합니다.
            self.d_optimizer.apply_gradients(
                zip(d_gradient, self.discriminator.trainable_variables)
            )

        # 생성자를 트레이닝합니다.
        # 잠재 벡터를 얻습니다.
        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
        with tf.GradientTape() as tape:
            # 생성기를 사용하여 가짜 이미지 생성
            generated_images = self.generator(random_latent_vectors, training=True)
            # 가짜 이미지에 대한 판별자 로짓을 얻습니다.
            gen_img_logits = self.discriminator(generated_images, training=True)
            # 생성자 손실을 계산합니다.
            g_loss = self.g_loss_fn(gen_img_logits)

        # 생성자 손실에 대한 그래디언트를 얻습니다.
        gen_gradient = tape.gradient(g_loss, self.generator.trainable_variables)
        # 생성자 옵티마이저를 사용하여 생성자의 가중치를 업데이트합니다.
        self.g_optimizer.apply_gradients(
            zip(gen_gradient, self.generator.trainable_variables)
        )
        return {"d_loss": d_loss, "g_loss": g_loss}

class GANMonitor(keras.callbacks.Callback):
    def __init__(self, num_img=6, latent_dim=128):
        self.num_img = num_img
        self.latent_dim = latent_dim

    def on_epoch_end(self, epoch, logs=None):
        random_latent_vectors = tf.random.normal(shape=(self.num_img, self.latent_dim))
        generated_images = self.model.generator(random_latent_vectors)
        generated_images = (generated_images * 127.5) + 127.5

        for i in range(self.num_img):
            img = generated_images[i].numpy()
            img = keras.utils.array_to_img(img)
            img.save("generated_img_{i}_{epoch}.png".format(i=i, epoch=epoch))

# 두 네트워크 모두에 대한 옵티마이저를 인스턴스화합니다.
# (learning_rate=0.0002, beta_1=0.5가 권장됩니다)
generator_optimizer = keras.optimizers.Adam(
    learning_rate=0.0002, beta_1=0.5, beta_2=0.9
)
discriminator_optimizer = keras.optimizers.Adam(
    learning_rate=0.0002, beta_1=0.5, beta_2=0.9
)


# 판별자에 대한 손실 함수를 정의합니다.
# 이는 (fake_loss - real_loss)여야 합니다.
# 나중에 이 손실 함수에 그래디언트 페널티를 추가합니다.
def discriminator_loss(real_img, fake_img):
    real_loss = tf.reduce_mean(real_img)
    fake_loss = tf.reduce_mean(fake_img)
    return fake_loss - real_loss


# 생성자의 손실 함수를 정의합니다.
def generator_loss(fake_img):
    return -tf.reduce_mean(fake_img)


# 트레이닝을 위한 에포크 수를 설정합니다.
epochs = 20

# customer `GANMonitor` Keras 콜백을 인스턴스화합니다.
cbk = GANMonitor(num_img=3, latent_dim=noise_dim)

# Wgan 모델을 얻습니다.
wgan = WGAN(
    discriminator=d_model,
    generator=g_model,
    latent_dim=noise_dim,
    discriminator_extra_steps=3,
)

# wgan 모델을 컴파일합니다.
wgan.compile(
    d_optimizer=discriminator_optimizer,
    g_optimizer=generator_optimizer,
    g_loss_fn=generator_loss,
    d_loss_fn=discriminator_loss,
)

# 트레이닝 시작
wgan.fit(train_images, batch_size=BATCH_SIZE, epochs=epochs, callbacks=[cbk])

from IPython.display import Image, display

display(Image("generated_img_0_19.png"))
display(Image("generated_img_1_19.png"))
display(Image("generated_img_2_19.png"))