{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Import libraries"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import warnings\nimport os\nimport time\nfrom glob import glob\nimport datetime\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport cv2\nimport xml\nimport xml.etree.ElementTree as ET\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport zipfile\nimport IPython.display as display\n\nimport tensorflow as tf\nfrom tensorflow.keras.layers import Input, Dense, Reshape, Flatten, Dropout, concatenate\nfrom tensorflow.keras.layers import BatchNormalization, Activation, ZeroPadding2D, LeakyReLU\nfrom tensorflow.keras.layers import Conv2D, Conv2DTranspose, UpSampling2D\nfrom tensorflow.keras.models import Sequential, Model\nfrom tensorflow.keras.initializers import TruncatedNormal, RandomNormal\nfrom tensorflow.keras.callbacks import LearningRateScheduler, EarlyStopping, History\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras import backend as K\n\n\ntf.enable_eager_execution()\nprint(tf.test.gpu_device_name())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(os.listdir(\"../input/\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# filepaths for image loading code\nROOT = '../input/'\n# list of all image file names in all-dogs\nIMAGES = os.listdir(ROOT + 'all-dogs/all-dogs')\n# list of all the annotation directories, each directory is a dog breed\nBREEDS = os.listdir(ROOT + 'annotation/Annotation/') \n\n# variables that determine how tensorflow will create batches after data load\nBUFFER_SIZE = 20000\nBATCH_SIZE = 32\n\n# weight initializers for the generator network\nWEIGHT_INIT = RandomNormal(mean=0.0, stddev=0.05)\n\n# generate/classify 64x64 images\nIMG_SIZE = 64\n\n# for training\nEPOCHS = 1000\nNOISE_SIZE = 128\nNB_EXAMPLES_TO_GENERATE = 16\n\n# for animated GIF\nseed = tf.random.normal([NB_EXAMPLES_TO_GENERATE, NOISE_SIZE])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load and preprocess images"},{"metadata":{"trusted":true},"cell_type":"code","source":"# dom = xml.dom.minidom.parse('../input/annotation/Annotation/n02097658-silky_terrier/n02097658_98') \n# pretty_xml_as_string = dom.toprettyxml()\n# print(pretty_xml_as_string)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Code slightly modified from user: cdeotte | https://www.kaggle.com/cdeotte/supervised-generative-dog-net\n\nimgs = []\nnames = []\n\n# CROP WITH BOUNDING BOXES TO GET DOGS ONLY\n# iterate through each directory in annotation\nfor breed in BREEDS:\n    # iterate through each file in the directory\n    for dog in os.listdir(ROOT+'annotation/Annotation/'+breed):\n        try: img = Image.open(ROOT+'all-dogs/all-dogs/'+dog+'.jpg')\n        except: continue\n        # Element Tree library allows for parsing xml and getting specific tag values\n        tree = ET.parse(ROOT+'annotation/Annotation/'+breed+'/'+dog)\n        # take a look at the print out of an xml previously to get what is going on\n        root = tree.getroot() # <annotation>\n        objects = root.findall('object') # <object>\n        for o in objects:\n            bndbox = o.find('bndbox') # <bndbox>\n            xmin = int(bndbox.find('xmin').text) # <xmin>\n            ymin = int(bndbox.find('ymin').text) # <ymin>\n            xmax = int(bndbox.find('xmax').text) # <xmax>\n            ymax = int(bndbox.find('ymax').text) # <ymax>\n            w = np.min((xmax - xmin, ymax - ymin))\n            img2 = img.crop((xmin, ymin, xmin+w, ymin+w))\n            img2 = img2.resize((IMG_SIZE, IMG_SIZE), Image.ANTIALIAS)\n            imgs.append(np.asarray(img2))\n            names.append(breed)\n\nimgs = np.array(imgs)\nnames[:] = map(str.lower, names)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"imgs.shape: {}\".format(imgs.shape))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# DISPLAY CROPPED IMAGES\n\n# list of 25 random numbers in range 0, idxIn\n# this allows for displaying random images in the for loop using x[k*5+j]\nx = np.random.randint(0, len(imgs), 25)\n\nfor k in range(5):\n    plt.figure(figsize=(15,3))\n    for j in range(5):\n        plt.subplot(1,5,j+1)\n        img = Image.fromarray(imgs[x[k*5+j], :, :, :].astype('uint8'))\n        plt.axis('off')\n        plt.title(names[x[k*5+j]].split('-')[1], fontsize=11)\n        plt.imshow(img)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# normalize the pixel values\nimgs = (imgs - 127.5) / 127.5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# view some images after normalization \n# plt.figure(figsize=(8,8))\n# for image in range(4):\n#     plt.subplot(2,2, image+1)\n#     plt.imshow((imgs[image]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# this is needed because the gradient functions from TF require float32 instead of float64\nimgs = imgs.astype(np.float32)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create tf dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"ds = tf.data.Dataset.from_tensor_slices(imgs).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\nprint(ds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Spectral norm"},{"metadata":{"trusted":true},"cell_type":"code","source":"def spectral_norm(w, iteration=1):\n   w_shape = w.shape.as_list()\n   w = tf.reshape(w, [-1, w_shape[-1]])\n\n   u = tf.get_variable(\"u\", [1, w_shape[-1]], initializer=tf.random_normal_initializer(), trainable=False)\n\n   u_hat = u\n   v_hat = None\n   for i in range(iteration):\n       \"\"\"\n       power iteration\n       Usually iteration = 1 will be enough\n       \"\"\"\n       v_ = tf.matmul(u_hat, tf.transpose(w))\n       v_hat = tf.nn.l2_normalize(v_)\n\n       u_ = tf.matmul(v_hat, w)\n       u_hat = tf.nn.l2_normalize(u_)\n\n   u_hat = tf.stop_gradient(u_hat)\n   v_hat = tf.stop_gradient(v_hat)\n\n   sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))\n\n   with tf.control_dependencies([u.assign(u_hat)]):\n       w_norm = w / sigma\n       w_norm = tf.reshape(w_norm, w_shape)\n\n   return w_norm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Pixelwise norm"},{"metadata":{"trusted":true},"cell_type":"code","source":"class PixelwiseNorm(tf.keras.layers.Layer):\n\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n\n\n    def call(self, x, eps=1e-8):\n        \"\"\"\n        :param x: input activations volume\n        :param alpha: small number for numerical stability\n        :return: y => pixel normalized activations\n        \"\"\"\n        y = tf.sqrt(tf.reduce_mean(x**2, axis=3, keepdims=True) + eps)\n        y = x / y  # normalize the input x volume\n        return y","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Adam with weight norm"},{"metadata":{"trusted":true},"cell_type":"code","source":"# adapted from keras.optimizers.Adam\nclass AdamWithWeightnorm(Adam):\n    def get_updates(self, loss, params):\n        grads = self.get_gradients(loss, params)\n        self.updates = [K.update_add(self.iterations, 1)]\n\n        lr = self.lr\n        if self.initial_decay > 0:\n            lr *= (1. / (1. + self.decay * K.cast(self.iterations, K.floatx())))\n\n        t = K.cast(self.iterations + 1, K.floatx())\n        lr_t = lr * K.sqrt(1. - K.pow(self.beta_2, t)) / (1. - K.pow(self.beta_1, t))\n\n        shapes = [K.get_variable_shape(p) for p in params]\n        ms = [K.zeros(shape) for shape in shapes]\n        vs = [K.zeros(shape) for shape in shapes]\n        self.weights = [self.iterations] + ms + vs\n\n        for p, g, m, v in zip(params, grads, ms, vs):\n\n            # if a weight tensor (len > 1) use weight normalized parameterization\n            # this is the only part changed w.r.t. keras.optimizers.Adam\n            ps = K.get_variable_shape(p)\n            if len(ps)>1:\n\n                # get weight normalization parameters\n                V, V_norm, V_scaler, g_param, grad_g, grad_V = get_weightnorm_params_and_grads(p, g)\n\n                # Adam containers for the 'g' parameter\n                V_scaler_shape = K.get_variable_shape(V_scaler)\n                m_g = K.zeros(V_scaler_shape)\n                v_g = K.zeros(V_scaler_shape)\n\n                # update g parameters\n                m_g_t = (self.beta_1 * m_g) + (1. - self.beta_1) * grad_g\n                v_g_t = (self.beta_2 * v_g) + (1. - self.beta_2) * K.square(grad_g)\n                new_g_param = g_param - lr_t * m_g_t / (K.sqrt(v_g_t) + self.epsilon)\n                self.updates.append(K.update(m_g, m_g_t))\n                self.updates.append(K.update(v_g, v_g_t))\n\n                # update V parameters\n                m_t = (self.beta_1 * m) + (1. - self.beta_1) * grad_V\n                v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(grad_V)\n                new_V_param = V - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)\n                self.updates.append(K.update(m, m_t))\n                self.updates.append(K.update(v, v_t))\n\n                # if there are constraints we apply them to V, not W\n                if getattr(p, 'constraint', None) is not None:\n                    new_V_param = p.constraint(new_V_param)\n\n                # wn param updates --> W updates\n                add_weightnorm_param_updates(self.updates, new_V_param, new_g_param, p, V_scaler)\n\n            else: # do optimization normally\n                m_t = (self.beta_1 * m) + (1. - self.beta_1) * g\n                v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(g)\n                p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)\n\n                self.updates.append(K.update(m, m_t))\n                self.updates.append(K.update(v, v_t))\n\n                new_p = p_t\n                # apply constraints\n                if getattr(p, 'constraint', None) is not None:\n                    new_p = p.constraint(new_p)\n                self.updates.append(K.update(p, new_p))\n        return self.updates\n\ndef get_weightnorm_params_and_grads(p, g):\n    ps = K.get_variable_shape(p)\n\n    # construct weight scaler: V_scaler = g/||V||\n    V_scaler_shape = (ps[-1],)  # assumes we're using tensorflow!\n    V_scaler = K.ones(V_scaler_shape)  # init to ones, so effective parameters don't change\n\n    # get V parameters = ||V||/g * W\n    norm_axes = [i for i in range(len(ps) - 1)]\n    V = p / tf.reshape(V_scaler, [1] * len(norm_axes) + [-1])\n\n    # split V_scaler into ||V|| and g parameters\n    V_norm = tf.sqrt(tf.reduce_sum(tf.square(V), norm_axes))\n    g_param = V_scaler * V_norm\n\n    # get grad in V,g parameters\n    grad_g = tf.reduce_sum(g * V, norm_axes) / V_norm\n    grad_V = tf.reshape(V_scaler, [1] * len(norm_axes) + [-1]) * \\\n             (g - tf.reshape(grad_g / V_norm, [1] * len(norm_axes) + [-1]) * V)\n\n    return V, V_norm, V_scaler, g_param, grad_g, grad_V\n\ndef add_weightnorm_param_updates(updates, new_V_param, new_g_param, W, V_scaler):\n    ps = K.get_variable_shape(new_V_param)\n    norm_axes = [i for i in range(len(ps) - 1)]\n\n    # update W and V_scaler\n    new_V_norm = tf.sqrt(tf.reduce_sum(tf.square(new_V_param), norm_axes))\n    new_V_scaler = new_g_param / new_V_norm\n    new_W = tf.reshape(new_V_scaler, [1] * len(norm_axes) + [-1]) * new_V_param\n    updates.append(K.update(W, new_W))\n    updates.append(K.update(V_scaler, new_V_scaler))\n\n# data based initialization for a given Keras model\ndef data_based_init(model, input):\n    # input can be dict, numpy array, or list of numpy arrays\n    if type(input) is dict:\n        feed_dict = input\n    elif type(input) is list:\n        feed_dict = {tf_inp: np_inp for tf_inp,np_inp in zip(model.inputs,input)}\n    else:\n        feed_dict = {model.inputs[0]: input}\n\n    # add learning phase if required\n    if model.uses_learning_phase and K.learning_phase() not in feed_dict:\n        feed_dict.update({K.learning_phase(): 1})\n\n    # get all layer name, output, weight, bias tuples\n    layer_output_weight_bias = []\n    for l in model.layers:\n        trainable_weights = l.trainable_weights\n        if len(trainable_weights) == 2:\n            W,b = trainable_weights\n            assert(l.built)\n            layer_output_weight_bias.append((l.name,l.get_output_at(0),W,b)) # if more than one node, only use the first\n\n    # iterate over our list and do data dependent init\n    sess = K.get_session()\n    for l,o,W,b in layer_output_weight_bias:\n        print('Performing data dependent initialization for layer ' + l)\n        m,v = tf.nn.moments(o, [i for i in range(len(o.get_shape())-1)])\n        s = tf.sqrt(v + 1e-10)\n        updates = tf.group(W.assign(W/tf.reshape(s,[1]*(len(W.get_shape())-1)+[-1])), b.assign((b-m)/s))\n        sess.run(updates, feed_dict)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generator model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# function will return a generator model\ndef make_generator():\n    model = Sequential([\n\n        Input(shape=(NOISE_SIZE,)),\n        # first layer with 32, 768 nodes expecting an input of vector size NOISE_SIZE (random noise)\n        Dense(4*4*64, use_bias=False),\n        # apply leaky relu activation: f(x) = {x if x > 0 : 0.01*x}\n        LeakyReLU(),\n        # Normalize the activations of the previous layer at each batch\n#         BatchNormalization(),\n        # reshape input to (8,8,512)\n        Reshape((4, 4, 64)),\n        # added later\n#         PixelwiseNorm(),\n\n        # 8 x 8\n        UpSampling2D((2, 2)),\n        Conv2D(128, (3, 3), padding='same', use_bias=False,\n                        kernel_initializer=WEIGHT_INIT),\n#         Conv2DTranspose(128, (3, 3), strides=(2,2), padding='same', use_bias=False,\n#                         kernel_initializer=WEIGHT_INIT),\n#                         kernel_regularizer=spectral_norm),\n        LeakyReLU(),\n#         BatchNormalization(),\n#         Dropout(0.1),\n#         PixelwiseNorm(),\n        \n\n        # 16 x 16\n        UpSampling2D((2, 2)),\n        Conv2D(128, (3, 3), padding='same', use_bias=False,\n                        kernel_initializer=WEIGHT_INIT),\n#         Conv2DTranspose(128, (3, 3), strides=(2,2), padding='same', use_bias=False,\n#                         kernel_initializer=WEIGHT_INIT),\n#                         kernel_regularizer=spectral_norm),\n        LeakyReLU(),\n#         BatchNormalization(),\n#         Dropout(0.1),\n#         PixelwiseNorm(),\n        \n        \n        # 32 x 32\n        UpSampling2D((2, 2)),\n        Conv2D(128, (3, 3), padding='same', use_bias=False,\n                        kernel_initializer=WEIGHT_INIT),\n#         Conv2DTranspose(128, (3, 3), strides=(2,2), padding='same', use_bias=False,\n#                         kernel_initializer=WEIGHT_INIT),\n#                         kernel_regularizer=spectral_norm),\n        LeakyReLU(),\n#         BatchNormalization(),\n#         PixelwiseNorm(),\n\n        # 64 x 64\n        UpSampling2D((2, 2)),\n        Conv2D(128, (3, 3), padding='same', use_bias=False,\n                        kernel_initializer=WEIGHT_INIT),\n#         Conv2DTranspose(128, (3, 3), strides=(2,2), padding='same', use_bias=False,\n#                         kernel_initializer=WEIGHT_INIT),\n#                         kernel_regularizer=spectral_norm),\n        LeakyReLU(),\n        \n        \n        Conv2D(3, (3, 3), padding='same', use_bias=False, activation='tanh', kernel_initializer=WEIGHT_INIT)\n        \n        \n#         Dense(3, activation='tanh', use_bias=False,\n#               kernel_initializer=WEIGHT_INIT),\n#               kernel_regularizer=spectral_norm)\n        \n#         Conv2D(3, (1, 1), padding='same', use_bias=False,\n#                kernel_initializer=WEIGHT_INIT)\n\n    ])\n    \n    return model\n\n\n# create an instance of the generator model defined\ngenerator = make_generator()\n\n# random noise vector\nnoise = tf.random.normal([1, NOISE_SIZE])\n\n# run the generator model with the noise vector as input\ngenerated_image = generator(noise, training=False)\n\nprint(generated_image.dtype)\n\n# display output\nplt.imshow(generated_image[0, :, :, :])\nprint(generated_image.shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Discriminator model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_discriminator():\n    model = tf.keras.Sequential([\n        \n        Input(shape=(IMG_SIZE, IMG_SIZE, 3)),\n        \n        # downsample to 32 x 32\n        Conv2D(128, (3, 3), strides=(2, 2), padding='same',\n               kernel_initializer=WEIGHT_INIT),\n        LeakyReLU(0.2),\n#         BatchNormalization(),\n        Dropout(0.25),\n    \n        # downsample to 16 x 16\n        Conv2D(128, (3, 3), strides=(2, 2), padding='same',\n               kernel_initializer=WEIGHT_INIT),\n#                kernel_regularizer=spectral_norm),\n        LeakyReLU(0.2),\n#         BatchNormalization(),\n        Dropout(0.25),\n\n        # downsample to 8 x 8\n        Conv2D(128, (3, 3), strides=(2, 2), padding='same',\n               kernel_initializer=WEIGHT_INIT),\n#                kernel_regularizer=spectral_norm),\n        LeakyReLU(0.2),\n#         BatchNormalization(),\n        Dropout(0.25),\n    \n        # downsample to 4 x 4\n        Conv2D(128, (3, 3), strides=(2, 2), padding='same',\n               kernel_initializer=WEIGHT_INIT),\n#                kernel_regularizer=spectral_norm),\n        LeakyReLU(0.2),\n#         BatchNormalization(),\n        Dropout(0.25),\n\n        # flatten input into 1-D and output a single a number from the last layer using sigmoid activation\n        Flatten(),\n        Dense(1, activation='sigmoid', kernel_initializer=WEIGHT_INIT)\n    ])\n\n    return model\n\n\ndiscriminator = make_discriminator()\n\ndecision = discriminator(generated_image)\nprint(decision)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss and optimizers"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Label smoothing -- technique from GAN hacks, instead of assigning 1/0 as class labels, we assign a random integer in range [0.7, 1.0] for positive class\n# and [0.0, 0.3] for negative class\n\ndef smooth_positive_labels(y):\n    return y - 0.3 + (np.random.random(y.shape) * 0.5)\n\n\ndef smooth_negative_labels(y):\n\treturn y + np.random.random(y.shape) * 0.3\n\n\n# Recomended to introduce some noise to the labels, so out of 1000 real labels, approximately 50 should be flipped to 0 (5%)\n# randomly flip some labels\n\ndef noisy_labels(y, p_flip=0.05):\n\t# determine the number of labels to flip\n\tn_select = int(p_flip * y.shape[0].value)\n\t# choose labels to flip\n\tflip_ix = choice([i for i in range(y.shape[0].value)], size=n_select)\n\t# invert the labels in place\n\ty[flip_ix] = 1 - y[flip_ix]\n\treturn y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# This method returns a helper function to compute cross entropy loss\n# code from tf dcgan tutorial\ncross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=False)\n\n# The Discriminator loss function\ndef discriminator_loss(real_output, fake_output):\n    real_output_smooth = smooth_positive_labels(tf.ones_like(real_output))\n#     real_output_noisy = noisy_labels(real_output_smooth, 0.05)\n    fake_output_smooth = smooth_negative_labels(tf.zeros_like(fake_output))\n#     fake_output_noisy = noisy_labels(fake_output_smooth, 0.05)\n    real_loss = cross_entropy(real_output_smooth, real_output)\n    fake_loss = cross_entropy(fake_output_smooth, fake_output)\n    total_loss = real_loss + fake_loss\n    return total_loss\n\n\n# The Generator loss function\ndef generator_loss(fake_output):\n    fake_output_smooth = smooth_negative_labels(tf.ones_like(fake_output))\n#     fake_output_noisy = noisy_labels(fake_output_smooth, 0.05)\n    return cross_entropy(fake_output_smooth, fake_output)\n\n\n# optimizers -- Adam\ngenerator_optimizer = AdamWithWeightnorm(lr=0.0002, beta_1=0.5)\ndiscriminator_optimizer = AdamWithWeightnorm(lr=0.00006, beta_1=0.5)\n# generator_optimizer = tf.train.AdamOptimizer(learning_rate=0.0002, beta1=0.5)\n# discriminator_optimizer = tf.train.AdamOptimizer(learning_rate=0.00006, beta1=0.5)\n# discriminator_optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.00006)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"# code from tf dcgan tutorial\ndef train_step(images, G_loss_list, D_loss_list):\n    noise = tf.random.normal([BATCH_SIZE, NOISE_SIZE])\n\n    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n        generated_images = generator(noise, training=True)\n        real_output = discriminator(images, training=True)\n        fake_output = discriminator(generated_images, training=True)\n        gen_loss = generator_loss(fake_output)\n        disc_loss = discriminator_loss(real_output, fake_output)\n\n    G_loss_list.append(gen_loss.numpy().mean())\n    D_loss_list.append(disc_loss.numpy().mean())\n    \n    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)\n    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n\n    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))\n    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(dataset, epochs, time_limit=32000):\n    training_start = time.time()\n\n    for epoch in range(epochs):\n        G_loss = []\n        D_loss = []\n\n        start = time.time()\n        for image_batch in dataset:\n            train_step(image_batch, G_loss, D_loss)\n        if (epoch % 10 == 0):\n#             display.clear_output(wait=True)\n            generate_and_save_images(generator,\n                                     epoch + 1,\n                                     seed)\n        print('epoch {:3d} - G loss: {:.4f} - D loss: {:.4f} - {:.2f} sec'.format(epoch + 1, G_loss[-1], D_loss[-1], time.time()-start))\n        if(time.time() - training_start > time_limit):\n            print(f\"Reached training time limit ({time_limit} s) at {time.time() - training_start:.2f}\")\n            break\n\n    # Generate after the final epoch\n    print(\"Final Epoch\")\n    generate_and_save_images(generator,\n                             epochs,\n                             seed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def generate_and_save_images(model, epoch, test_input):\n    # Notice `training` is set to False.\n    # This is so all layers run in inference mode (batchnorm).\n    predictions = model(test_input, training=False)\n    fig = plt.figure(figsize=(8,8))\n    for i in range(predictions.shape[0]):\n        plt.subplot(4, 4, i+1)\n        plt.imshow((predictions[i, :, :, :] + 1.) / 2.)\n        plt.axis('off')\n    plt.savefig(\"img_at_epoch_{}.png\".format(epoch+1))\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nprint('Starting training')\ntrain(ds, EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Save submission"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\n# SAVE TO ZIP FILE NAMED IMAGES.ZIP\nz = zipfile.PyZipFile('images.zip', mode='w')\n\nfilename = 'gen_model.h5'\ntf.keras.models.save_model(\n    generator,\n    filename,\n    overwrite=True,\n    include_optimizer=True,\n    save_format=None\n)\n\nfor k in range(10000):\n    generated_image = generator(tf.random.normal([1, NOISE_SIZE]), training=False)\n    f = str(k) + '.png'\n    img = ((generated_image[0,:,:,:] + 1.) / 2.).numpy()\n    tf.keras.preprocessing.image.save_img(\n        f,\n        img,\n        scale=True\n    )\n    z.write(f)\n    os.remove(f)\n\nz.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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":1}