{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom matplotlib import pyplot as plt\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\nfrom tensorflow.keras import layers, Sequential, Model\nfrom IPython.display import clear_output\nfrom tqdm import tqdm\nimport time\nAUTOTUNE = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# HYPER PARAMETERS\n\nWIDTH, HEIGHT = 256, 256\nEPOCHS = 20\nBUFFER_SIZE = 400\nBATCH_SIZE = 1\nOUTPUT_CHANNEL = 3\n\n# LOSSES\nLAYER_COUNT = 0\nLAMBDA = 100","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH_X = \"/kaggle/input/gan-getting-started/photo_jpg\"\nPATH_Y = \"/kaggle/input/gan-getting-started/monet_jpg\"\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def crop(img):\n    crop_img = tf.image.random_crop(img, size=[HEIGHT, WIDTH, 3])\n    return crop_img\n\n\ndef train_preprocess(img_path):\n\n    img = tf.io.read_file(img_path)\n    img = tf.image.decode_jpeg(img)    \n    img = tf.image.resize(img, [HEIGHT+20, WIDTH+20], method = tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n    img = crop(img)\n    \n    # normalising\n    img = tf.cast(img, tf.float32)\n    img = img/127.5 -1\n    \n    # random flipping    \n    img = tf.image.random_flip_left_right(img)\n\n    return img\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_x = tf.data.Dataset.list_files(PATH_X + '/*.jpg').map(train_preprocess).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\ntrain_y = tf.data.Dataset.list_files(PATH_Y + '/*.jpg').map(train_preprocess).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_x = next(iter(train_x))[0]\nsample_x = tf.reshape(sample_x, [-1, 256, 256, 3])\nsample_y = next(iter(train_y))[0]\nsample_y = tf.reshape(sample_y, [-1, 256, 256, 3])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.subplot(121)\nplt.imshow((sample_x[0]+1.)/2.)\nplt.subplot(122)\nplt.imshow((sample_y[0]+1.)/2.)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def add_layer(filters, kernel_size,batchnorm = True):\n    init = tf.random_normal_initializer(0., 0.02)\n    blocks = Sequential()\n    blocks.add(layers.Conv2D(filters, kernel_size=kernel_size, strides = 4, kernel_initializer=init,padding='same', use_bias=False))\n    if batchnorm:\n        blocks.add(layers.BatchNormalization())\n    blocks.add(layers.LeakyReLU())\n    return blocks\n\ndef add_trans_layer(filters, kernel_size, dropout=True):\n    init = tf.random_normal_initializer(0., 0.02)\n    blocks = Sequential()\n    blocks.add(layers.Conv2DTranspose(filters, kernel_size, strides=4, kernel_initializer=init, use_bias=False, padding='same'))\n    blocks.add(layers.BatchNormalization())\n    if dropout:\n        blocks.add(layers.Dropout(0.4))\n    blocks.add(layers.LeakyReLU())\n    return blocks","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef make_gen():\n    inputs = layers.Input(shape = [256, 256, 3])\n\n    down = [\n            add_layer(64, 5, False),\n            add_layer(128, 5),\n            add_layer(256, 5),\n            add_layer(512, 5)\n    ]\n\n    up = [\n            add_trans_layer(256, 5),\n            add_trans_layer(128, 5),\n            add_trans_layer(64, 5, dropout=False)\n    ]\n\n    init = tf.random_normal_initializer(0., 0.02)\n\n    output = layers.Conv2DTranspose(OUTPUT_CHANNEL, kernel_size=5, strides=(4,4), padding='same', kernel_initializer=init, activation='tanh')\n\n    x = inputs\n\n    stack = []\n\n    for d in down:\n        x = d(x)\n        stack.append(x)\n    \n    stack.pop(-1)\n\n    for a in up:\n        x = a(x)\n        x = layers.Concatenate()([x, stack[-1]])\n        stack.pop(-1)\n\n    x = output(x)\n\n    return Model(inputs = inputs, outputs=x)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Gx = make_gen()\nGy = make_gen()\ntf.keras.utils.plot_model(Gx, show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def gen_loss(fake_output):\n    return cross_entropy(tf.ones_like(fake_output), fake_output)\n\nGx_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\nGy_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_disc():\n    init = tf.random_normal_initializer(0., 0.05)\n\n    input_shape = [256, 256, 3]\n\n    down = [\n            add_layer(64, (5,5), batchnorm=False),\n            add_layer(256, (5,5)),\n            layers.ZeroPadding2D(),\n            layers.Conv2D(512, kernel_size=(5,5), kernel_initializer=init, use_bias=False),\n            layers.BatchNormalization(),\n            layers.LeakyReLU(),\n            layers.ZeroPadding2D(),\n            layers.Conv2D(1, kernel_size=(5,5), kernel_initializer=init),\n            layers.Flatten(),\n            layers.Dense(1, activation='tanh')\n    ]\n\n    inp = layers.Input(shape=input_shape)\n\n    x = inp\n    for layer in down:\n        x = layer(x)\n  \n    return Model(inputs=inp, outputs = x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Dx = make_disc()\nDy = make_disc()\n\ntf.keras.utils.plot_model(Dx, show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def disc_loss(real_output, fake_output):\n    real_loss = cross_entropy(tf.ones_like(real_output), real_output)\n    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)\n    return fake_loss + real_loss\n\nDx_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)\nDy_opt = tf.keras.optimizers.Adam(1e-5, beta_1=0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def cyclic_loss(real_img, cycled_img):\n    return LAMBDA*tf.reduce_mean(abs(real_img-cycled_img))\n\ndef identity_loss(real_img, same_img):\n    return 0.5*LAMBDA*tf.reduce_mean(abs(real_img-same_img))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@tf.function\ndef train_step(real_x, real_y):\n    with tf.GradientTape(persistent=True) as tape:\n        fake_y = Gx(real_x, training=True)\n        cycled_x = Gy(fake_y, training=True)\n        fake_x = Gy(real_y, training=True)\n        cycled_y = Gx(fake_x, training=True)\n        same_x = Gx(real_x, training=True)\n        same_y = Gy(real_y, training=True)\n\n        Dx_real = Dx(real_x, training=True)\n        Dx_fake = Dx(fake_x, training=True)\n\n        Dy_real = Dy(real_y, training=True)\n        Dy_fake = Dy(fake_y, training=True)\n\n        Dx_loss = disc_loss(Dx_real, Dx_fake)\n        Dy_loss = disc_loss(Dy_real, Dy_fake)\n        \n        cycled_loss = cyclic_loss(real_x, cycled_x) + cyclic_loss(real_y, cycled_y)\n        \n        Gx_loss = gen_loss(Dy_fake) + identity_loss(real_x, same_x) + cycled_loss \n        Gy_loss = gen_loss(Dx_fake) + identity_loss(real_y, same_y) + cycled_loss\n\n\n    Gx_grad = tape.gradient(Gx_loss,Gx.trainable_variables)\n    Gx_opt.apply_gradients(zip(Gx_grad, Gx.trainable_variables))\n\n    Gy_grad = tape.gradient(Gy_loss, Gy.trainable_variables)\n    Gy_opt.apply_gradients(zip(Gy_grad, Gy.trainable_variables))\n\n    Dx_grad = tape.gradient(Dx_loss, Dx.trainable_variables)\n    Dx_opt.apply_gradients(zip(Dx_grad, Dx.trainable_variables))\n\n    Dy_grad = tape.gradient(Dy_loss, Dy.trainable_variables)\n    Dy_opt.apply_gradients(zip(Dy_grad, Dy.trainable_variables))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(train_x, train_y, EPOCHS):\n    init_time = time.time()\n    for epoch in range(EPOCHS):\n        start = time.time()\n        for x_batch, y_batch in zip(train_x, train_y):\n            train_step(x_batch, y_batch)\n        clear_output(wait=True)\n        print('Training...')\n        print(f\"|{(epoch+1)*'#' + (EPOCHS-epoch-1)*'-'}|\")\n        print(f'Time taken on Epoch {epoch+1} : {time.time()-start}')\n        print(f'Total time ellapsed so far : {time.time()-init_time}')\n    clear_output(wait=True)\n    print(\"Model Trained.\")\n    print(f'Total time ellapsed : {time.time()-init_time}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train(train_x, train_y, EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure()\nplt.subplot(121)\nplt.imshow(sample_x[0])\nplt.subplot(122)\nplt.imshow(Gx(sample_x)[0])","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}