title cyclegan aknainhttps twitter comaknain date created 20200812 last modified 20200812 description implementation of cyclegan accelerator gpu cyclegan cyclegan is a model that aims to solve the imagetoimage translation problem the goal of the imagetoimage translation problem is to learn the mapping between an input image and an output image using a training set of aligned image pairs however obtaining paired examples isn t always feasible cyclegan tries to learn this mapping without requiring paired inputoutput images using cycleconsistent adversarial networks paperhttps arxiv orgabs1703 10593 original implementationhttps github comjunyanzpytorchcycleganandpix2pix setup prepare the dataset in this example we will be using the horse to zebrahttps www tensorflow orgdatasetscatalogcyclegancycleganhorse2zebra dataset load the horsezebra dataset using tensorflowdatasets define the standard image size size of the random crops to be used during training weights initializer for the layers gamma initializer for instance normalization map values in the range 1 1 random flip resize to the original size first random crop to 256x256 normalize the pixel values in the range 1 1 only resizing and normalization for the test images create dataset objects apply the preprocessing operations to the training data apply the preprocessing operations to the test data visualize some samples building blocks used in the cyclegan generators and discriminators implements reflection padding as a layer args paddingtuple amount of padding for the spatial dimensions returns a padded tensor with the same type as the input tensor build the generators the generator consists of downsampling blocks nine residual blocks and upsampling blocks the structure of the generator is the following c7s164 conv block with relu activation filter size of 7 d128 2 downsampling blocks d256 r256 r256 r256 r256 r256 9 residual blocks r256 r256 r256 r256 u128 2 upsampling blocks u64 c7s13 last conv block with tanh activation filter size of 7 downsampling residual blocks upsampling final block build the discriminators the discriminators implement the following architecture c64c128c256c512 get the generators get the discriminators build the cyclegan model we will override the trainstep method of the model class for training via fit x is horse and y is zebra for cyclegan we need to calculate different kinds of losses for the generators and discriminators we will perform the following steps here 1 pass real images through the generators and get the generated images 2 pass the generated images back to the generators to check if we can predict the original image from the generated image 3 do an identity mapping of the real images using the generators 4 pass the generated images in 1 to the corresponding discriminators 5 calculate the generators total loss adversarial cycle identity 6 calculate the discriminators loss 7 update the weights of the generators 8 update the weights of the discriminators 9 return the losses in a dictionary horse to fake zebra zebra to fake horse y2x cycle horse to fake zebra to fake horse x y x cycle zebra to fake horse to fake zebra y x y identity mapping discriminator output generator adversarial loss generator cycle loss generator identity loss total generator loss discriminator loss get the gradients for the generators get the gradients for the discriminators update the weights of the generators update the weights of the discriminators create a callback that periodically saves generated images a callback to generate and save images after each epoch def initself numimg4 self numimg numimg def onepochendself epoch logsnone ax plt subplots4 2 figsize12 12 for i img in enumeratetesthorses takeself numimg prediction self model gengimg0 numpy prediction prediction 127 5 127 5 astypenp uint8 img img0 127 5 127 5 numpy astypenp uint8 axi 0 imshowimg axi 1 imshowprediction axi 0 settitleinput image axi 1 settitletranslated image axi 0 axisoff axi 1 axisoff prediction keras utils arraytoimgprediction prediction save generatedimgiepoch png formatii epochepoch 1 plt show plt close train the endtoend model loss function for evaluating adversarial loss advlossfn keras losses meansquarederror define the loss function for the generators def generatorlossfnfake fakeloss advlossfnops oneslikefake fake return fakeloss define the loss function for the discriminators def discriminatorlossfnreal fake realloss advlossfnops oneslikereal real fakeloss advlossfnops zeroslikefake fake return realloss fakeloss 0 5 create cycle gan model cycleganmodel cyclegan generatorggeng generatorfgenf discriminatorxdiscx discriminatorydiscy compile the model cycleganmodel compile gengoptimizerkeras optimizers adamlearningrate2e4 beta10 5 genfoptimizerkeras optimizers adamlearningrate2e4 beta10 5 discxoptimizerkeras optimizers adamlearningrate2e4 beta10 5 discyoptimizerkeras optimizers adamlearningrate2e4 beta10 5 genlossfngeneratorlossfn disclossfndiscriminatorlossfn callbacks plotter ganmonitor checkpointfilepath cyclegancheckpoints epoch 03d weights h5 modelcheckpointcallback keras callbacks modelcheckpoint filepathcheckpointfilepath saveweightsonlytrue here we will train the model for just one epoch as each epoch takes around 7 minutes on a single p100 backed machine cycleganmodel fit tf data dataset ziptrainhorses trainzebras epochs1 callbacksplotter modelcheckpointcallback test the performance of the model this model was trained for 90 epochs we will be loading those weights here once the weights are loaded we will take a few samples from the test data and check the model s performance load the checkpoints weightfile cyclegankeras2 0modelcheckpointscyclegancheckpoints 090 weights h5 cycleganmodel loadweightsweightfile printweights loaded successfully ax plt subplots4 2 figsize10 15 for i img in enumeratetesthorses take4 prediction cycleganmodel geng predictimg0 prediction prediction 127 5 127 5 astypenp uint8 img img0 127 5 127 5 numpy astypenp uint8 axi 0 imshowimg axi 1 imshowprediction axi 0 settitleinput image axi 0 settitleinput image axi 1 settitletranslated image axi 0 axisoff axi 1 axisoff prediction keras utils arraytoimgprediction prediction savepredictedimgi png formatii plt tightlayout plt show cyclegan cyclegan is a model that aims to solve the image to image translation problem the goal of the image to image translation problem is to learn the mapping between an input image and an output image using a training set of aligned image pairs however obtaining paired examples isn t always feasible cyclegan tries to learn this mapping without requiring paired input output images using cycle consistent adversarial networks paper https arxiv org abs 1703 10593 original implementation https github com junyanz pytorch cyclegan and pix2pix setup prepare the dataset in this example we will be using the horse to zebra https www tensorflow org datasets catalog cycle_gan cycle_ganhorse2zebra dataset load the horse zebra dataset using tensorflow datasets define the standard image size size of the random crops to be used during training weights initializer for the layers gamma initializer for instance normalization map values in the range 1 1 random flip resize to the original size first random crop to 256x256 normalize the pixel values in the range 1 1 only resizing and normalization for the test images create dataset objects apply the preprocessing operations to the training data apply the preprocessing operations to the test data visualize some samples building blocks used in the cyclegan generators and discriminators implements reflection padding as a layer args padding tuple amount of padding for the spatial dimensions returns a padded tensor with the same type as the input tensor build the generators the generator consists of downsampling blocks nine residual blocks and upsampling blocks the structure of the generator is the following c7s1 64 conv block with relu activation filter size of 7 d128 2 downsampling blocks d256 r256 r256 r256 r256 r256 9 residual blocks r256 r256 r256 r256 u128 2 upsampling blocks u64 c7s1 3 last conv block with tanh activation filter size of 7 downsampling residual blocks upsampling final block build the discriminators the discriminators implement the following architecture c64 c128 c256 c512 get the generators get the discriminators build the cyclegan model we will override the train_step method of the model class for training via fit x is horse and y is zebra for cyclegan we need to calculate different kinds of losses for the generators and discriminators we will perform the following steps here 1 pass real images through the generators and get the generated images 2 pass the generated images back to the generators to check if we can predict the original image from the generated image 3 do an identity mapping of the real images using the generators 4 pass the generated images in 1 to the corresponding discriminators 5 calculate the generators total loss adversarial cycle identity 6 calculate the discriminators loss 7 update the weights of the generators 8 update the weights of the discriminators 9 return the losses in a dictionary horse to fake zebra zebra to fake horse y2x cycle horse to fake zebra to fake horse x y x cycle zebra to fake horse to fake zebra y x y identity mapping discriminator output generator adversarial loss generator cycle loss generator identity loss total generator loss discriminator loss get the gradients for the generators get the gradients for the discriminators update the weights of the generators update the weights of the discriminators create a callback that periodically saves generated images a callback to generate and save images after each epoch train the end to end model loss function for evaluating adversarial loss define the loss function for the generators define the loss function for the discriminators create cycle gan model compile the model callbacks here we will train the model for just one epoch as each epoch takes around 7 minutes on a single p100 backed machine test the performance of the model this model was trained for 90 epochs we will be loading those weights here once the weights are loaded we will take a few samples from the test data and check the model s performance shell curl lo https github com freedomtan cyclegan keras archive refs tags 2 0 zip unzip qq 2 0 zip load the checkpoints
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import numpy as np
import matplotlib.pyplot as plt
import keras
from keras import layers
from keras import ops

import tensorflow as tf
import tensorflow_datasets as tfds

tfds.disable_progress_bar()
autotune = tf.data.AUTOTUNE




dataset, _ = tfds.load(
    "cycle_gan/horse2zebra", with_info=True, as_supervised=True
)
train_horses, train_zebras = dataset["trainA"], dataset["trainB"]
test_horses, test_zebras = dataset["testA"], dataset["testB"]


orig_img_size = (286, 286)

input_img_size = (256, 256, 3)

kernel_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)

gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)

buffer_size = 256
batch_size = 1


def normalize_img(img):
    img = ops.cast(img, dtype="float32")
    
    return (img / 127.5) - 1.0


def preprocess_train_image(img, label):
    
    img = tf.image.random_flip_left_right(img)
    
    img = tf.image.resize(img, [*orig_img_size])
    
    img = tf.image.random_crop(img, size=[*input_img_size])
    
    img = normalize_img(img)
    return img


def preprocess_test_image(img, label):
    
    img = tf.image.resize(img, [input_img_size[0], input_img_size[1]])
    img = normalize_img(img)
    return img






train_horses = (
    train_horses.map(preprocess_train_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)
train_zebras = (
    train_zebras.map(preprocess_train_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)


test_horses = (
    test_horses.map(preprocess_test_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)
test_zebras = (
    test_zebras.map(preprocess_test_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)





_, ax = plt.subplots(4, 2, figsize=(10, 15))
for i, samples in enumerate(zip(train_horses.take(4), train_zebras.take(4))):
    horse = (((samples[0][0] * 127.5) + 127.5).numpy()).astype(np.uint8)
    zebra = (((samples[1][0] * 127.5) + 127.5).numpy()).astype(np.uint8)
    ax[i, 0].imshow(horse)
    ax[i, 1].imshow(zebra)
plt.show()





class ReflectionPadding2D(layers.Layer):
    

    def __init__(self, padding=(1, 1), **kwargs):
        self.padding = tuple(padding)
        super().__init__(**kwargs)

    def call(self, input_tensor, mask=None):
        padding_width, padding_height = self.padding
        padding_tensor = [
            [0, 0],
            [padding_height, padding_height],
            [padding_width, padding_width],
            [0, 0],
        ]
        return ops.pad(
            input_tensor, ops.convert_to_tensor(padding_tensor), mode="reflect"
        )


def residual_block(
    x,
    activation,
    kernel_initializer=kernel_init,
    kernel_size=(3, 3),
    strides=(1, 1),
    padding="valid",
    gamma_initializer=gamma_init,
    use_bias=False,
):
    dim = x.shape[-1]
    input_tensor = x

    x = ReflectionPadding2D()(input_tensor)
    x = layers.Conv2D(
        dim,
        kernel_size,
        strides=strides,
        kernel_initializer=kernel_initializer,
        padding=padding,
        use_bias=use_bias,
    )(x)
    x = layers.GroupNormalization(
        groups=x.shape[-1], gamma_initializer=gamma_initializer
    )(x)
    x = activation(x)

    x = ReflectionPadding2D()(x)
    x = layers.Conv2D(
        dim,
        kernel_size,
        strides=strides,
        kernel_initializer=kernel_initializer,
        padding=padding,
        use_bias=use_bias,
    )(x)
    x = layers.GroupNormalization(
        groups=x.shape[-1], gamma_initializer=gamma_initializer
    )(x)
    x = layers.add([input_tensor, x])
    return x


def downsample(
    x,
    filters,
    activation,
    kernel_initializer=kernel_init,
    kernel_size=(3, 3),
    strides=(2, 2),
    padding="same",
    gamma_initializer=gamma_init,
    use_bias=False,
):
    x = layers.Conv2D(
        filters,
        kernel_size,
        strides=strides,
        kernel_initializer=kernel_initializer,
        padding=padding,
        use_bias=use_bias,
    )(x)
    x = layers.GroupNormalization(
        groups=x.shape[-1], gamma_initializer=gamma_initializer
    )(x)
    if activation:
        x = activation(x)
    return x


def upsample(
    x,
    filters,
    activation,
    kernel_size=(3, 3),
    strides=(2, 2),
    padding="same",
    kernel_initializer=kernel_init,
    gamma_initializer=gamma_init,
    use_bias=False,
):
    x = layers.Conv2DTranspose(
        filters,
        kernel_size,
        strides=strides,
        padding=padding,
        kernel_initializer=kernel_initializer,
        use_bias=use_bias,
    )(x)
    x = layers.GroupNormalization(
        groups=x.shape[-1], gamma_initializer=gamma_initializer
    )(x)
    if activation:
        x = activation(x)
    return x





def get_resnet_generator(
    filters=64,
    num_downsampling_blocks=2,
    num_residual_blocks=9,
    num_upsample_blocks=2,
    gamma_initializer=gamma_init,
    name=None,
):
    img_input = layers.Input(shape=input_img_size, name=name + "_img_input")
    x = ReflectionPadding2D(padding=(3, 3))(img_input)
    x = layers.Conv2D(
        filters, (7, 7), kernel_initializer=kernel_init, use_bias=False
    )(x)
    x = layers.GroupNormalization(
        groups=x.shape[-1], gamma_initializer=gamma_initializer
    )(x)
    x = layers.Activation("relu")(x)

    
    for _ in range(num_downsampling_blocks):
        filters *= 2
        x = downsample(x, filters=filters, activation=layers.Activation("relu"))

    
    for _ in range(num_residual_blocks):
        x = residual_block(x, activation=layers.Activation("relu"))

    
    for _ in range(num_upsample_blocks):
        filters //= 2
        x = upsample(x, filters, activation=layers.Activation("relu"))

    
    x = ReflectionPadding2D(padding=(3, 3))(x)
    x = layers.Conv2D(3, (7, 7), padding="valid")(x)
    x = layers.Activation("tanh")(x)

    model = keras.models.Model(img_input, x, name=name)
    return model





def get_discriminator(
    filters=64, kernel_initializer=kernel_init, num_downsampling=3, name=None
):
    img_input = layers.Input(shape=input_img_size, name=name + "_img_input")
    x = layers.Conv2D(
        filters,
        (4, 4),
        strides=(2, 2),
        padding="same",
        kernel_initializer=kernel_initializer,
    )(img_input)
    x = layers.LeakyReLU(0.2)(x)

    num_filters = filters
    for num_downsample_block in range(3):
        num_filters *= 2
        if num_downsample_block < 2:
            x = downsample(
                x,
                filters=num_filters,
                activation=layers.LeakyReLU(0.2),
                kernel_size=(4, 4),
                strides=(2, 2),
            )
        else:
            x = downsample(
                x,
                filters=num_filters,
                activation=layers.LeakyReLU(0.2),
                kernel_size=(4, 4),
                strides=(1, 1),
            )

    x = layers.Conv2D(
        1,
        (4, 4),
        strides=(1, 1),
        padding="same",
        kernel_initializer=kernel_initializer,
    )(x)

    model = keras.models.Model(inputs=img_input, outputs=x, name=name)
    return model



gen_G = get_resnet_generator(name="generator_G")
gen_F = get_resnet_generator(name="generator_F")


disc_X = get_discriminator(name="discriminator_X")
disc_Y = get_discriminator(name="discriminator_Y")





class CycleGan(keras.Model):
    def __init__(
        self,
        generator_G,
        generator_F,
        discriminator_X,
        discriminator_Y,
        lambda_cycle=10.0,
        lambda_identity=0.5,
    ):
        super().__init__()
        self.gen_G = generator_G
        self.gen_F = generator_F
        self.disc_X = discriminator_X
        self.disc_Y = discriminator_Y
        self.lambda_cycle = lambda_cycle
        self.lambda_identity = lambda_identity

    def call(self, inputs):
        return (
            self.disc_X(inputs),
            self.disc_Y(inputs),
            self.gen_G(inputs),
            self.gen_F(inputs),
        )

    def compile(
        self,
        gen_G_optimizer,
        gen_F_optimizer,
        disc_X_optimizer,
        disc_Y_optimizer,
        gen_loss_fn,
        disc_loss_fn,
    ):
        super().compile()
        self.gen_G_optimizer = gen_G_optimizer
        self.gen_F_optimizer = gen_F_optimizer
        self.disc_X_optimizer = disc_X_optimizer
        self.disc_Y_optimizer = disc_Y_optimizer
        self.generator_loss_fn = gen_loss_fn
        self.discriminator_loss_fn = disc_loss_fn
        self.cycle_loss_fn = keras.losses.MeanAbsoluteError()
        self.identity_loss_fn = keras.losses.MeanAbsoluteError()

    def train_step(self, batch_data):
        
        real_x, real_y = batch_data

        
        
        
        
        
        
        
        
        
        
        
        
        
        

        with tf.GradientTape(persistent=True) as tape:
            
            fake_y = self.gen_G(real_x, training=True)
            
            fake_x = self.gen_F(real_y, training=True)

            
            cycled_x = self.gen_F(fake_y, training=True)
            
            cycled_y = self.gen_G(fake_x, training=True)

            
            same_x = self.gen_F(real_x, training=True)
            same_y = self.gen_G(real_y, training=True)

            
            disc_real_x = self.disc_X(real_x, training=True)
            disc_fake_x = self.disc_X(fake_x, training=True)

            disc_real_y = self.disc_Y(real_y, training=True)
            disc_fake_y = self.disc_Y(fake_y, training=True)

            
            gen_G_loss = self.generator_loss_fn(disc_fake_y)
            gen_F_loss = self.generator_loss_fn(disc_fake_x)

            
            cycle_loss_G = (
                self.cycle_loss_fn(real_y, cycled_y) * self.lambda_cycle
            )
            cycle_loss_F = (
                self.cycle_loss_fn(real_x, cycled_x) * self.lambda_cycle
            )

            
            id_loss_G = (
                self.identity_loss_fn(real_y, same_y)
                * self.lambda_cycle
                * self.lambda_identity
            )
            id_loss_F = (
                self.identity_loss_fn(real_x, same_x)
                * self.lambda_cycle
                * self.lambda_identity
            )

            
            total_loss_G = gen_G_loss + cycle_loss_G + id_loss_G
            total_loss_F = gen_F_loss + cycle_loss_F + id_loss_F

            
            disc_X_loss = self.discriminator_loss_fn(disc_real_x, disc_fake_x)
            disc_Y_loss = self.discriminator_loss_fn(disc_real_y, disc_fake_y)

        
        grads_G = tape.gradient(total_loss_G, self.gen_G.trainable_variables)
        grads_F = tape.gradient(total_loss_F, self.gen_F.trainable_variables)

        
        disc_X_grads = tape.gradient(
            disc_X_loss, self.disc_X.trainable_variables
        )
        disc_Y_grads = tape.gradient(
            disc_Y_loss, self.disc_Y.trainable_variables
        )

        
        self.gen_G_optimizer.apply_gradients(
            zip(grads_G, self.gen_G.trainable_variables)
        )
        self.gen_F_optimizer.apply_gradients(
            zip(grads_F, self.gen_F.trainable_variables)
        )

        
        self.disc_X_optimizer.apply_gradients(
            zip(disc_X_grads, self.disc_X.trainable_variables)
        )
        self.disc_Y_optimizer.apply_gradients(
            zip(disc_Y_grads, self.disc_Y.trainable_variables)
        )

        return {
            "G_loss": total_loss_G,
            "F_loss": total_loss_F,
            "D_X_loss": disc_X_loss,
            "D_Y_loss": disc_Y_loss,
        }





class GANMonitor(keras.callbacks.Callback):
    

    def __init__(self, num_img=4):
        self.num_img = num_img

    def on_epoch_end(self, epoch, logs=None):
        _, ax = plt.subplots(4, 2, figsize=(12, 12))
        for i, img in enumerate(test_horses.take(self.num_img)):
            prediction = self.model.gen_G(img)[0].numpy()
            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)
            img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)

            ax[i, 0].imshow(img)
            ax[i, 1].imshow(prediction)
            ax[i, 0].set_title("Input image")
            ax[i, 1].set_title("Translated image")
            ax[i, 0].axis("off")
            ax[i, 1].axis("off")

            prediction = keras.utils.array_to_img(prediction)
            prediction.save(
                "generated_img_{i}_{epoch}.png".format(i=i, epoch=epoch + 1)
            )
        plt.show()
        plt.close()






adv_loss_fn = keras.losses.MeanSquaredError()




def generator_loss_fn(fake):
    fake_loss = adv_loss_fn(ops.ones_like(fake), fake)
    return fake_loss



def discriminator_loss_fn(real, fake):
    real_loss = adv_loss_fn(ops.ones_like(real), real)
    fake_loss = adv_loss_fn(ops.zeros_like(fake), fake)
    return (real_loss + fake_loss) * 0.5



cycle_gan_model = CycleGan(
    generator_G=gen_G,
    generator_F=gen_F,
    discriminator_X=disc_X,
    discriminator_Y=disc_Y,
)


cycle_gan_model.compile(
    gen_G_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
    gen_F_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
    disc_X_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
    disc_Y_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
    gen_loss_fn=generator_loss_fn,
    disc_loss_fn=discriminator_loss_fn,
)

plotter = GANMonitor()
checkpoint_filepath = (
    "cyclegan_checkpoints.{epoch:03d}.weights.h5"
)
model_checkpoint_callback = keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath, save_weights_only=True
)



cycle_gan_model.fit(
    tf.data.Dataset.zip((train_horses, train_zebras)),
    epochs=1,
    callbacks=[plotter, model_checkpoint_callback],
)












weight_file = "./cyclegan-keras-2.0/model_checkpoints/cyclegan_checkpoints.090.weights.h5"
cycle_gan_model.load_weights(weight_file)
print("Weights loaded successfully")

_, ax = plt.subplots(4, 2, figsize=(10, 15))
for i, img in enumerate(test_horses.take(4)):
    prediction = cycle_gan_model.gen_G.predict(img)[0]
    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)
    img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)

    ax[i, 0].imshow(img)
    ax[i, 1].imshow(prediction)
    ax[i, 0].set_title("Input image")
    ax[i, 0].set_title("Input image")
    ax[i, 1].set_title("Translated image")
    ax[i, 0].axis("off")
    ax[i, 1].axis("off")

    prediction = keras.utils.array_to_img(prediction)
    prediction.save("predicted_img_{i}.png".format(i=i))
plt.tight_layout()
plt.show()