{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nfrom matplotlib import pyplot as plt\ntry:\n tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n print('Device:', tpu.master())\n tf.config.experimental_connect_to_cluster(tpu)\n tf.tpu.experimental.initialize_tpu_system(tpu)\n strategy = tf.distribute.experimental.TPUStrategy(tpu)\nexcept:\n strategy = tf.distribute.get_strategy()\nprint('Number of replicas:', strategy.num_replicas_in_sync)\n\n\n# !/opt/bin/nvidia-smi","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_data():\n \"\"\"thank for Amy Jang for this piece of code : https://www.kaggle.com/amyjang/monet-cyclegan-tutorial#Load-in-the-data \n \"\"\"\n from kaggle_datasets import KaggleDatasets\n GCS_PATH = KaggleDatasets().get_gcs_path()\n MONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\n\n PHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\n\n\n AUTOTUNE = tf.data.experimental.AUTOTUNE\n IMAGE_SIZE = [256, 256]\n\n def decode_image(image):\n image = tf.image.decode_jpeg(image, channels=3)\n image = (tf.cast(image, tf.float32) / 127.5) - 1\n image = tf.reshape(image, [*IMAGE_SIZE, 3])\n return image\n\n def read_tfrecord(example):\n tfrecord_format = {\n \"image_name\": tf.io.FixedLenFeature([], tf.string),\n \"image\": tf.io.FixedLenFeature([], tf.string),\n \"target\": tf.io.FixedLenFeature([], tf.string)\n }\n example = tf.io.parse_single_example(example, tfrecord_format)\n image = decode_image(example['image'])\n return image\n\n\n def load_dataset(filenames, labeled=True, ordered=False):\n dataset = tf.data.TFRecordDataset(filenames)\n dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n return dataset\n\n\n monet_ds = load_dataset(MONET_FILENAMES, labeled=True)\n photo_ds = load_dataset(PHOTO_FILENAMES, labeled=True)\n return monet_ds, photo_ds","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# A liitle tutorial for `Conv2DTranspose()`\nFor using Deep Convolutional GAN, it is important to determine the strides of transposed convolution so that it outputs the desired image size, i.e. (256, 256) in this example. Firstly, recall how strides affect the output of common convolution: if padding compensates the reduced size due to kernel (i.e. $p_{h}=k_{h}-1 \\text { and } p_{w}=k_{w}-1$) and the input height and width are divisible by the strides on the height and width, then $$ \\left(n_{h} / s_{h}\\right) \\times\\left(n_{w} / s_{w}\\right)$$\n\nSee [more details](https://d2l.ai/chapter_convolutional-neural-networks/padding-and-strides.html) and the following code cell could be used for experiment.\n\n
\n \n
From Up-sampling with Transposed Convolution
\n
","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Here padding='same' would eliminate the affect of kernel on the output size, \n# by changing strides from 1 to 2, we can see the size halved: 4=>2.\nx = tf.random.normal((2, 4, 4, 1))\nlayers.Conv2D(64, (3, 3), padding='same', strides=(1,1))(x).shape[1:3]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x = tf.random.normal((2, 256, 256, 1))\nlayers.Conv2D(64, (3, 3), padding='same', strides=(4,4))(x).shape[1:3]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"To calculate the size of output for transposed convolution, we multiply the size by the strides. In this example, we want to upsample for size 256 after 3 `Conv2DTranspose()` with strides 1 for first and strides 2 for last two. At the begining of the first `Conv2DTranspose`, we should get input size: 256/(2*2)=64","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Build Generator","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class Generator(keras.Model):\n \n def __init__(self):\n super(Generator, self).__init__()\n \n # z: [b, 100] => [b, 64, 64, 256] => [b, 256, 256, 3]\n self.fc = layers.Dense(64*64*256, use_bias=False)\n #layers.BatchNormalization()\n \n\n self.conv1 = layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False)\n #assert model.output_shape == (None, 7, 7, 128)\n self.bn1 = layers.BatchNormalization()\n\n\n self.conv2 = layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False)\n #assert model.output_shape == (None, 14, 14, 64)\n self.bn2 = layers.BatchNormalization()\n\n\n self.conv3 = layers.Conv2DTranspose(3, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')\n #assert model.output_shape == (None, 28, 28, 1)\n\n\n \n def call(self, inputs, training=None):\n x = self.fc(inputs)\n x = tf.reshape(x, (-1, 64, 64, 256))\n x = tf.nn.leaky_relu(x)\n \n x = tf.nn.leaky_relu(self.bn1(self.conv1(x), training=training))\n x = tf.nn.leaky_relu(self.bn2(self.conv2(x), training=training))\n x = tf.tanh(self.conv3(x))\n \n return x\n \n ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Build Discriminator","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class Discriminator(keras.Model):\n \n def __init__(self):\n super(Discriminator, self).__init__()\n \n # input: [b, 256, 256, 3] => output: [b, 1]\n \n self.conv1 = layers.Conv2D(64, (5, 5), strides=(3, 3), padding='valid')\n #layers.Dropout(0.3)\n \n self.conv2 = layers.Conv2D(128, (5, 5), strides= (3, 3), padding='valid')\n self.bn2 = layers.BatchNormalization()\n #layers.Dropout(0.3)\n \n self.conv3 = layers.Conv2D(256, (5, 5), strides= (3, 3), padding='valid')\n self.bn3 = layers.BatchNormalization()\n \n # [b, h, w, 3] => [b, -1]\n self.flatten = layers.Flatten()\n self.fc = layers.Dense(1)\n\n\n \n def call(self, inputs, training=None):\n \n x = tf.nn.leaky_relu(self.conv1(inputs))\n x = tf.nn.leaky_relu(self.bn2(self.conv2(x), training=training))\n x = tf.nn.leaky_relu(self.bn3(self.conv3(x), training=training))\n \n x = self.flatten(x)\n logits = self.fc(x)\n \n return logits","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define the loss","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n def celoss_ones(logits):\n # logits shape: [b, 1]\n # labels shape: [b] = [1] * num_pics\n loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,\n labels=tf.ones_like(logits))\n\n return tf.reduce_mean(loss)\nwith strategy.scope():\n def celoss_zeros(logits):\n # logits shape: [b, 1]\n # labels shape: [b] = [0] * num_pics\n loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,\n labels=tf.zeros_like(logits))\n\n return tf.reduce_mean(loss)\n\n\n \n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class DeepConvGan(keras.Model):\n \n def __init__(self, generator, discriminator):\n \n super(DeepConvGan, self).__init__()\n self.generator = generator\n self.discriminator = discriminator\n \n \n def compile(self, gen_optimizer, \n disc_optimizer,\n celoss_ones, \n celoss_zeros):\n \n super(DeepConvGan, self).compile()\n self.gen_optimizer = gen_optimizer\n self.disc_optimizer = disc_optimizer\n self.celoss_ones = celoss_ones\n self.celoss_zeros = celoss_zeros\n \n \n def train_step(self, batch_data):\n \n batch_x, batch_z = batch_data\n \n # train disc\n with tf.GradientTape() as tape:\n # random noise to fake monet \n fake_image = self.generator(batch_z, training=True)\n \n # disc the fake and real image\n d_fake_logits = self.discriminator(fake_image, training=True)\n d_real_logits = self.discriminator(batch_x, training=True)\n \n # get discriminator loss\n d_loss_real = self.celoss_ones(d_real_logits)\n d_loss_fake = self.celoss_zeros(d_fake_logits)\n monet_disc_loss = d_loss_fake + d_loss_real\n \n grads = tape.gradient(monet_disc_loss, self.discriminator.trainable_variables)\n \n self.disc_optimizer.apply_gradients(zip(grads, self.discriminator.trainable_variables))\n\n \n # Train generator\n with tf.GradientTape() as tape:\n # get generator loss\n fake_image = self.generator(batch_z, training=True)\n d_fake_logits = self.discriminator(fake_image, training=True)\n photo_gen_loss = self.celoss_ones(d_fake_logits)\n \n grads = tape.gradient(photo_gen_loss, self.generator.trainable_variables)\n self.gen_optimizer.apply_gradients(zip(grads, self.generator.trainable_variables))\n \n return {\n \"photo_gen_loss\": photo_gen_loss,\n \"monet_disc_loss\": monet_disc_loss\n }\n\n ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Main Code cell \ntf.random.set_seed(22)\nnp.random.seed(22)\n\n# set up hyperparameters\nz_dim = 100\nepochs = 300\nbatch_size = 128\nlearning_rate = 2e-4\n\nwith strategy.scope():\n gen_optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)\n disc_optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)\n\n# load data and train\nmonet_ds, photo_ds = load_data()\nmonet_ds = monet_ds.batch(batch_size).repeat()\nphoto_ds = photo_ds.batch(batch_size).repeat()\nmonet_iter = iter(monet_ds)\nphoto_iter = iter(photo_ds)\n\nbatch_z = tf.random.uniform([batch_size, z_dim], minval=-1., maxval=1.)\ngen_ds = tf.data.Dataset.from_tensor_slices(batch_z).batch(batch_size).repeat()\n \nwith strategy.scope():\n \n \n # initialize instances\n discriminator = Discriminator()\n discriminator.build(input_shape=(None, 256, 256, 3))\n generator = Generator()\n generator.build(input_shape=(None, z_dim))\n\n \n dc_gan = DeepConvGan(generator, discriminator)\n \n dc_gan.compile(gen_optimizer=gen_optimizer, \n disc_optimizer=disc_optimizer,\n celoss_ones=celoss_ones, \n celoss_zeros=celoss_zeros)\n \n \n# gen_loss_fn=gen_loss_fn, \n# disc_loss_fn=gen_loss_fn\n \ndc_gan.fit(tf.data.Dataset.zip((monet_ds, gen_ds)),epochs=300, steps_per_epoch=600 // batch_size)\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img = dc_gan.generator(tf.random.uniform([1, z_dim], minval=-1., maxval=1.), training=False).numpy()[0,:,:,:]\nplt.imshow(img*0.5+0.5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Compare to real monet paintings, I know this generated one is bad (just like noisy colorful dots, but still arts!!!). I consider reducing batch size. But after first round, the loss tends to become NaN. I am not sure what happens here. In my sense, when I reduce the batch size, the competing game between discriminator and generator should become more significant. Anyway, welcome to comment for the mistakes. Since I just do this for weekend's entertainment and have not studied nuances of the GAN story, Would come back after more understanding.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.imshow(next(monet_iter).numpy()[0]*0.5+0.5)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}