gan-tutorial-3-how-to-use-tpus-to-train-a-gan.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.8.16","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"source":"<a href=\"https://www.kaggle.com/code/peremartramanonellas/gan-tutorial-3-how-to-use-tpus-to-train-a-gan?scriptVersionId=136918892\" target=\"_blank\"><img align=\"left\" alt=\"Kaggle\" title=\"Open in Kaggle\" src=\"https://kaggle.com/static/images/open-in-kaggle.svg\"></a>","metadata":{},"cell_type":"markdown"},{"cell_type":"markdown","source":"\n","metadata":{"execution":{"iopub.status.busy":"2023-03-07T22:35:25.738271Z","iopub.execute_input":"2023-03-07T22:35:25.740194Z","iopub.status.idle":"2023-03-07T22:35:37.795955Z","shell.execute_reply.started":"2023-03-07T22:35:25.740112Z","shell.execute_reply":"2023-03-07T22:35:37.794371Z"}}},{"cell_type":"code","source":"from IPython.core.display import HTML\nHTML(\"\"\"\n<style>\n.output_png {\n    display: table-cell;\n    text-align: center;\n    vertical-align: middle;\n    horizontal-align: middle;\n}\nh1 {\n    text-align: center;\n    background-color: #6bacf5;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n    style=\"font-family:verdana;\"\n}\n\nh2 {\n    text-align: center;\n    background-color: #83c2ff;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n}\n\nh3 {\n    text-align: center;\n    background-color: pink;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n}\n\nh4 {\n    text-align: center;\n    background-color: pink;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n}\n\nbody, p {\n    font-family: monospace;\n    font-size: 18px;\n    color: charcoal;\n}\ndiv {\n    font-size: 14px;\n    margin: 0;\n\n}\n\n\n</style>\n\"\"\")","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-07-16T09:02:41.917086Z","iopub.execute_input":"2023-07-16T09:02:41.917606Z","iopub.status.idle":"2023-07-16T09:02:41.941542Z","shell.execute_reply.started":"2023-07-16T09:02:41.917577Z","shell.execute_reply":"2023-07-16T09:02:41.940873Z"},"trusted":true},"execution_count":1,"outputs":[{"execution_count":1,"output_type":"execute_result","data":{"text/plain":"<IPython.core.display.HTML object>","text/html":"\n<style>\n.output_png {\n    display: table-cell;\n    text-align: center;\n    vertical-align: middle;\n    horizontal-align: middle;\n}\nh1 {\n    text-align: center;\n    background-color: #6bacf5;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n    style=\"font-family:verdana;\"\n}\n\nh2 {\n    text-align: center;\n    background-color: #83c2ff;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n}\n\nh3 {\n    text-align: center;\n    background-color: pink;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n}\n\nh4 {\n    text-align: center;\n    background-color: pink;\n    padding: 10px;\n    margin: 0;\n    font-family: monospace;\n    color:DimGray;\n    border-radius: 2px\n}\n\nbody, p {\n    font-family: monospace;\n    font-size: 18px;\n    color: charcoal;\n}\ndiv {\n    font-size: 14px;\n    margin: 0;\n\n}\n\n\n</style>\n"},"metadata":{}}]},{"cell_type":"markdown","source":"\n# Using TPUs and distributed training to accelerate our Notebooks. \n\nLooking to speed up your GAN training? In this noptebook, I'll try to show you how to utilize TPUs in Kaggle or Google Colab to drastically reduce training time. \n\nFollow this  step-by-step guide and take your GANs to the next level!\n\nThe code in this notebook wortks in Kaggle and in Colab. If you can't acces to TPUs at Kaggle, just try with the same notebook at [Google Colab](https://colab.research.google.com/drive/1p6sQqiu4kWeDpxu91C0MQBX9P6qSwmPG?usp=sharing)\n\n### Feel Free to fork or edit the noteboook for you own convenience. Please consider UPVOTING IT. It helps others to discover the notebook, and it encourages me to continue publishing.\n\nThis is the third notebook in the GANs tutorial. Here is a link to the other Notebooks. \nhttps://www.kaggle.com/code/peremartramanonellas/gan-tutorial-first-dcgan-using-tensorflow\nhttps://www.kaggle.com/code/peremartramanonellas/gan-tutorial-2-generating-color-images\n\nI'm going to use the CelebA Dataset. It's a Dataset composed of 200,000 images with famous people's faces. \n\n### **If you want more explanations than the ones in this Notebook you can check the full tutorial at medium:** \nhttps://medium.com/@peremartra/list/gans-from-zero-to-hero-d8e6cb773f93\n\n\n\n","metadata":{}},{"cell_type":"markdown","source":"## Importing Libraries\n","metadata":{}},{"cell_type":"code","source":"#Model face detection library, we can use the face detection in the notebooks. \n!pip install mtcnn","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:34:33.73021Z","iopub.execute_input":"2023-03-09T22:34:33.730616Z","iopub.status.idle":"2023-03-09T22:34:42.606266Z","shell.execute_reply.started":"2023-03-09T22:34:33.730572Z","shell.execute_reply":"2023-03-09T22:34:42.60535Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#import libraries\nimport tensorflow as tf\nfrom tensorflow.keras.utils import plot_model\nfrom tensorflow.keras import layers\nimport tensorflow.keras as keras\nfrom IPython import display\n\nimport os\nimport zipfile\nimport glob\nimport urllib.request\nfrom enum import Enum\nfrom tqdm import tqdm\nfrom functools import partial\n\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nfrom PIL import Image\nfrom IPython.display import display\nfrom IPython.display import Image as IpyImage\nimport imageio\nimport cv2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-09T22:34:42.60826Z","iopub.execute_input":"2023-03-09T22:34:42.608562Z","iopub.status.idle":"2023-03-09T22:34:50.506292Z","shell.execute_reply.started":"2023-03-09T22:34:42.608534Z","shell.execute_reply":"2023-03-09T22:34:50.505145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Face detection library\nimport mtcnn\nprint(mtcnn.__version__)\n     ","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:34:50.507835Z","iopub.execute_input":"2023-03-09T22:34:50.508464Z","iopub.status.idle":"2023-03-09T22:34:50.517257Z","shell.execute_reply.started":"2023-03-09T22:34:50.50843Z","shell.execute_reply":"2023-03-09T22:34:50.516165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Support function to print the images. \ndef plot_results(images, n_cols=None, unnorm=True, save_path = \"\" ):\n\n    n_cols = n_cols or len(images)\n    n_rows = (len(images) - 1) // n_cols + 1\n\n    if images.shape[-1] == 1:\n        images = np.squeeze(images, axis=-1)\n\n    plt.figure(figsize=(12, 12))\n    \n    for index, image in enumerate(images):\n        plt.subplot(n_rows, n_cols, index + 1)\n        imageorg = image\n\n        #We can print normalized and unnormalized images, just using the unnorm parameter. \n        if unnorm:\n          imageorg = (image +1) * 0.5\n          \n        plt.imshow(imageorg)\n        plt.axis(\"off\")\n\n        if save_path:\n          plt.savefig(save_path)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:34:50.518425Z","iopub.execute_input":"2023-03-09T22:34:50.518735Z","iopub.status.idle":"2023-03-09T22:34:50.531114Z","shell.execute_reply.started":"2023-03-09T22:34:50.518699Z","shell.execute_reply":"2023-03-09T22:34:50.529739Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Activating the TPU\nThis code aims to establish execution strategy. \n\nThe first thing is to connect to a TPU. Once connected, it's time to create the strategy with *tf.distribute.TPUStrategy*. \n\n\nIndicating, that we are going to execute the code in a distributed way. \n","metadata":{}},{"cell_type":"code","source":"try: \n  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\nexcept ValueError: \n  raise BaseException(\"CAN'T CONNECT TO A TPU\")\n\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nstrategy = tf.distribute.TPUStrategy(tpu)","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:34:50.534624Z","iopub.execute_input":"2023-03-09T22:34:50.535487Z","iopub.status.idle":"2023-03-09T22:34:56.281119Z","shell.execute_reply.started":"2023-03-09T22:34:50.535444Z","shell.execute_reply":"2023-03-09T22:34:56.280028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\nI'm going to use the Dataset FacesA, getting it from a storage in google, but there are more sources with the Dataset. (Just in case mine disappears). \nGoogle directory: https://storage.googleapis.com/learning-datasets/Resources/archive.zip\nKaggle: https://www.kaggle.com/datasets/jessicali9530/celeba-dataset\nMmlab: https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html\nTensorFlow: https://www.tensorflow.org/datasets/catalog/celeb_a?hl=es-419\n\nThe first thing to do is create the directory. \nAfter that, download the images. \nAnd finally, unzip the file. \n\nThe images will be in the directory: \n/tmp/celeb/img_align_celeba/img_align_celeba","metadata":{}},{"cell_type":"code","source":"# make a data directory\ntry:\n  os.mkdir('/tmp/celeb')\nexcept OSError:\n  pass\n\n# download the dataset archive\ndata_url = \"https://storage.googleapis.com/learning-datasets/Resources/archive.zip\"\ndata_file_name = \"archive.zip\"\ndownload_dir = '/tmp/celeb/'\nurllib.request.urlretrieve(data_url, data_file_name)\n\n# extract the zipped file\nzip_ref = zipfile.ZipFile(data_file_name, 'r')\nzip_ref.extractall(download_dir)\nzip_ref.close()","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:34:56.282343Z","iopub.execute_input":"2023-03-09T22:34:56.282719Z","iopub.status.idle":"2023-03-09T22:35:50.192821Z","shell.execute_reply.started":"2023-03-09T22:34:56.282683Z","shell.execute_reply":"2023-03-09T22:35:50.191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Time to load and transform the images. I have two functions: \n\n* load_faces: Cuts the central area of the image, then adjust the size. \n* load_faces_MTCNN: Uses the MTCNN library to get the coordinates where the face is image.","metadata":{}},{"cell_type":"code","source":"#Function to load the faces. \n#Crop & center the images because the faces is almost always in the center. \ndef load_faces(image_paths, resize, max_images):\n  crop_size = 128\n  if (max_images == 0): \n    max_images =  len(image_paths)\n    print(max_images)\n\n  images = np.zeros((max_images, resize, resize, 3), np.uint8)\n\n  for i, path in tqdm(enumerate(image_paths)):\n    with Image.open(path) as img:\n      left = (img.size[0] - crop_size) // 2\n      top = (img.size[1] - crop_size) // 2\n      right = left + crop_size\n      bottom = top + crop_size\n      img = img.crop((left, top, right, bottom))\n      img = img.resize((resize, resize), Image.LANCZOS)\n      images[i] = np.asarray(img, np.uint8)\n\n      if (i >= max_images-1):\n        break\n\n  return images","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:35:50.195682Z","iopub.execute_input":"2023-03-09T22:35:50.196865Z","iopub.status.idle":"2023-03-09T22:35:50.208228Z","shell.execute_reply.started":"2023-03-09T22:35:50.196828Z","shell.execute_reply":"2023-03-09T22:35:50.205357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Funtion to load the faces. Use the MTCNN lybrary to detect where the face is. \ndef load_faces_MTCNN(image_paths, resize, max_images):\n  MTCNN_model = mtcnn.MTCNN()\n  if (max_images == 0): \n    max_images =  len(image_paths)\n    print(max_images)\n\n  images = np.zeros((max_images, resize, resize, 3), np.uint8)\n  for i, path in enumerate(image_paths):\n    with Image.open(path) as img:\n\n      img = img.convert('RGB')\n      #img_pixels = np.asarray(img)\n      \n      face = MTCNN_model.detect_faces(img_pixels)\n      if len(face) == 0:\n        continue \n      x1, y1, width, height = face[0]['box']\n      x1, y1 = abs(x1), abs(y1)\n      x2, y2 = x1 + width, y1 + height\n\n      img = img.crop((x1, y1, x2, y2))\n      img = img.resize((resize, resize), Image.LANCZOS)\n      images[i] = np.asarray(img, np.uint8)\n      if (i >= max_images-1):\n        break\n  return images","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:35:50.209585Z","iopub.execute_input":"2023-03-09T22:35:50.209912Z","iopub.status.idle":"2023-03-09T22:35:50.27919Z","shell.execute_reply.started":"2023-03-09T22:35:50.20988Z","shell.execute_reply":"2023-03-09T22:35:50.277767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This is the function responsible for loading the Dataset.  I'm, calling load_faces. However, it is possible to change the load_faces function call to a load_faces_MTCNN call, and use the MTCNN library to detect where the faces are. \n\nEach of the images is passed through the preprocess function, to normalizing the pixels to values between -1 and 1.","metadata":{}},{"cell_type":"code","source":"#Load the pictures in the dataset. You can indicate the max_images. I recommend \n#that for testing you use 1000 images, and when you want to see the result final \n#use the max number of images, indicating 0 to max_images. \ndef load_celeba(batch_size, resize=80, max_images=0):\n  \"\"\"Creates batches of preprocessed images from the JPG files\n  Args:\n    batch_size - batch size\n    resize - size in pixels to resize the images\n    crop_size - size to crop from the image\n  \n  Returns:\n    prepared dataset\n  \"\"\"\n\n  # initialize zero-filled array equal to the size of the dataset\n  image_paths = sorted(glob.glob(\"/tmp/celeb/img_align_celeba/img_align_celeba/*.jpg\"))\n  \n  print(\"Creating Images\")\n\n  # crop and resize the raw images then put into the array\n  #choose wich function you want to use. \n  images = load_faces(image_paths, resize, max_images)\n  #images = load_faces_MTCNN(image_paths, resize, max_images)\n  \n  #Plot the 5 first images. \n  plot_results(images[0:5], unnorm=False)\n\n  # split the images array into two\n  split_n = images.shape[0] // 2\n  images1, images2 = images[:split_n], images[split_n:2 * split_n]\n  del images\n\n  # preprocessing function to convert the pixel values into the range [-1,1]\n  #Is a GAN Hack to normalize the pixels of the images \n  def preprocess(img):\n      x = tf.cast(img, tf.float32) / 127.5 - 1.0\n      return x\n  \n  # use the preprocessing function on the arrays and create batches\n  dataset = tf.data.Dataset.from_tensor_slices((images1, images2))\n  dataset = dataset.map(\n      lambda x1, x2: (preprocess(x1), preprocess(x2))\n  ).shuffle(4096).batch(batch_size, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)\n  \n  return dataset","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:35:50.28087Z","iopub.execute_input":"2023-03-09T22:35:50.281212Z","iopub.status.idle":"2023-03-09T22:35:50.293316Z","shell.execute_reply.started":"2023-03-09T22:35:50.281178Z","shell.execute_reply":"2023-03-09T22:35:50.29185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It's possible to adjust the maximum number of images to load. Although I recommend a minimum of 1000. In case we want to use all the images in the dataset, we only have to give a value of 0 to the variable max_images.","metadata":{}},{"cell_type":"code","source":"# use the function above to load and prepare the dataset\n#Note how the batch_size is multiplied by strategy.num_replicas_in_sync\nprint(strategy.num_replicas_in_sync)\nbatch_size = 8\nbatch_size = batch_size * strategy.num_replicas_in_sync\ndataset = load_celeba(batch_size, max_images=10000)\nout_dir = \"celeba_out\"","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:35:50.294974Z","iopub.execute_input":"2023-03-09T22:35:50.295881Z","iopub.status.idle":"2023-03-09T22:36:03.617495Z","shell.execute_reply.started":"2023-03-09T22:35:50.295841Z","shell.execute_reply":"2023-03-09T22:36:03.616597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"As we are using a distributed strategy to execute the notebook the Dataset must be prepared for that. ","metadata":{}},{"cell_type":"code","source":"    # create distributed dataset\n    dataset = strategy.experimental_distribute_dataset(dataset)\n    \n    # set the loss function\n    loss_func = tf.keras.losses.BinaryCrossentropy(\n        from_logits=True, \n        reduction=tf.keras.losses.Reduction.NONE\n    )","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:03.618638Z","iopub.execute_input":"2023-03-09T22:36:03.619199Z","iopub.status.idle":"2023-03-09T22:36:03.767989Z","shell.execute_reply.started":"2023-03-09T22:36:03.619168Z","shell.execute_reply":"2023-03-09T22:36:03.767067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generator & Discriminator\n\nI have a function to create Generators, and other to create Discriminators. They are really easy to use, and it is easy to adapt the Generator and Discriminator to different shape images. \n\nAs you know a GAN is composed of two models, the Generator and the Discriminator. The Generator creates images. \n\nThe discriminator tries to figure out if an image belongs to the original Dataset, or is a Generated one. ","metadata":{}},{"cell_type":"code","source":"def adapt_generator(initial_0, nodes, upsamplings, multnodes = 1.0, endnodes = 3, input_noise=100):\n  #initial_0, initial_1: size of the initial mini image. \n  #nodes: nodes in the first Dense layers. \n  #upsamplings: number og upsamplings bucles. \n  #multnodes: a multiplicator to modify the nodes in each upsampling bucle. \n  #endnodes: nodes of the last layer. \n  #input_noise: size of the noise. \n\n  model = keras.models.Sequential()\n\n  #First Dense layer. \n  model.add(keras.Input(shape=(1, 1, 128)))\n\n  #Upsampling bucles. \n  nodeslayers = nodes\n  model.add(keras.layers.Conv2DTranspose(nodeslayers , kernel_size=initial_0, strides=1, padding=\"valid\", \n                                 use_bias=False))\n\n  \n  for i in range(upsamplings-1):\n    nodeslayers = int(nodeslayers * multnodes)\n    model.add(keras.layers.Conv2DTranspose(nodeslayers , kernel_size=4, strides=2, padding=\"SAME\", \n                                 activation=keras.layers.LeakyReLU(alpha=0.2)))\n    model.add(keras.layers.BatchNormalization())\n\n  #last upsample and last layer. \n  model.add(keras.layers.Conv2DTranspose(endnodes, kernel_size=4, strides=2, padding=\"SAME\", \n                                 activation='tanh'))\n  \n  return model","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:03.770085Z","iopub.execute_input":"2023-03-09T22:36:03.770471Z","iopub.status.idle":"2023-03-09T22:36:03.778727Z","shell.execute_reply.started":"2023-03-09T22:36:03.770438Z","shell.execute_reply":"2023-03-09T22:36:03.777635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def adapt_discriminator(nodes, downsamples, multnodes = 1.0, in_shape=[32, 32, 3]):\n  #nodes: nodes in the first Dense layers.\n  #downsamples: number of downsamples bucles. \n  #multnodes: a multiplicator to modify the nodes in each downsample bucle.\n  #in_shape: Shape of the input image. \n\n  model = keras.models.Sequential()\n  \n  #input layer % first downsample \n  model.add(keras.layers.Conv2D(nodes, kernel_size=5, strides=2, padding=\"SAME\",\n                        activation=keras.layers.LeakyReLU(0.2),\n                        input_shape=in_shape))\n  model.add(keras.layers.Dropout(0.4))\n  \n  #creating downsamples\n  nodeslayers = nodes\n  for i in range(downsamples - 1):\n    nodeslayers = int(nodeslayers * multnodes)\n    model.add(keras.layers.Conv2D(nodeslayers, kernel_size=3, strides=2, padding=\"SAME\",\n                        activation=keras.layers.LeakyReLU(0.2)))\n    model.add(keras.layers.Dropout(0.4))\n  \n  #ending model\n  model.add(keras.layers.Flatten())\n  model.add(keras.layers.Dense(1, activation=\"sigmoid\"))\n  \n  return model","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:03.781342Z","iopub.execute_input":"2023-03-09T22:36:03.781722Z","iopub.status.idle":"2023-03-09T22:36:03.790037Z","shell.execute_reply.started":"2023-03-09T22:36:03.781691Z","shell.execute_reply":"2023-03-09T22:36:03.789134Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Settings\nresize = 80\nshape = (resize, resize, 3)\n\n# Build the GAN\nwith strategy.scope():\n    # create the generator model\n    model_G = adapt_generator(5, nodes=128, upsamplings=4, multnodes=1, endnodes=3, input_noise=100)\n\n    # create the discriminator model\n    model_D = adapt_discriminator(128, 5, multnodes=1, in_shape=shape)\n    \n    # print summaries\n    model_G.summary()\n    model_D.summary()\n\n    # set optimizers\n    param_G = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)\n    param_D = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:03.793555Z","iopub.execute_input":"2023-03-09T22:36:03.795012Z","iopub.status.idle":"2023-03-09T22:36:06.907568Z","shell.execute_reply.started":"2023-03-09T22:36:03.794947Z","shell.execute_reply":"2023-03-09T22:36:06.906834Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(model_G, show_shapes=True, \n           show_layer_names=True)","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:06.908489Z","iopub.execute_input":"2023-03-09T22:36:06.909227Z","iopub.status.idle":"2023-03-09T22:36:07.162035Z","shell.execute_reply.started":"2023-03-09T22:36:06.9092Z","shell.execute_reply":"2023-03-09T22:36:07.160855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(model_D, show_shapes=True, \n           show_layer_names=True)","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:07.163468Z","iopub.execute_input":"2023-03-09T22:36:07.163787Z","iopub.status.idle":"2023-03-09T22:36:07.297692Z","shell.execute_reply.started":"2023-03-09T22:36:07.163746Z","shell.execute_reply":"2023-03-09T22:36:07.296113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The distributed function is used to decorate the train function of the GAN. \n\nThe train function will be executed in multiple instances and we need to treat the return of the function to join the different values returned, we need to get the values like it were returned by a simple execution.\n\nWe can SUM (with reduction) or concatenate the variables returned. ","metadata":{}},{"cell_type":"code","source":"class Reduction(Enum):\n    SUM = 0\n    CONCAT = 1\n#TThis decorated function indicates how to concatenate the values \n#returned by all the functions working in the different distributed \n#TPU's. \n#We have two possibilites. return a reducted SUM of each process, or \n# a concatenation. \ndef distributed(*reduction_flags):\n    def _decorator(fun):\n        def per_replica_reduction(z, flag):\n            if flag == Reduction.SUM:\n                return strategy.reduce(tf.distribute.ReduceOp.SUM, z, axis=None)\n            elif flag == Reduction.CONCAT:\n                z_list = strategy.experimental_local_results(z)\n                return tf.concat(z_list, axis=0)\n            else:\n                raise NotImplementedError()\n\n        @tf.function\n        def _decorated_fun(*args, **kwargs):\n            fun_result = strategy.run(fun, args=args, kwargs=kwargs)\n\n            assert type(fun_result) is tuple\n            return tuple((per_replica_reduction(fr, rf) for fr, rf in zip(fun_result, reduction_flags)))\n        return _decorated_fun\n    return _decorator","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:07.299357Z","iopub.execute_input":"2023-03-09T22:36:07.29967Z","iopub.status.idle":"2023-03-09T22:36:07.309054Z","shell.execute_reply.started":"2023-03-09T22:36:07.299641Z","shell.execute_reply":"2023-03-09T22:36:07.308229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Our training function is decorated with the function distributed. As we are returning 3 variables, we must indicate how to concatenate all of them. \n\nIn the call to @distributed we are indicating the Sum of the first two parameters, while a concatenation is applied to the third. It makes sense because the first two are numbers that show the loss of the discriminator and the generator, and the third is a list of the images generated.\n","metadata":{}},{"cell_type":"code","source":"@distributed(Reduction.SUM, Reduction.SUM, Reduction.CONCAT)\ndef train_on_batch(real_img1, real_img2):\n    '''trains the GAN on a given batch'''\n    # concatenate the real image inputs\n    real_img = tf.concat([real_img1, real_img2], axis=0)\n\n    # PHASE ONE - train the discriminator\n    with tf.GradientTape() as d_tape:\n\n        # create noise input\n        z = tf.random.normal(shape=(real_img.shape[0], 1, 1, z_dim))\n\n        # generate fake images\n        fake_img = model_G(z)\n\n        # feed the fake images to the discriminator\n        fake_out = model_D(fake_img)\n\n        # feed the real images to the discriminator\n        real_out = model_D(real_img)\n\n        # use the loss function to measure how well the discriminator\n        # labels fake or real images\n        d_fake_loss = loss_func(tf.zeros_like(fake_out), fake_out)\n        d_real_loss = loss_func(tf.ones_like(real_out), real_out)\n\n        # get the total loss\n        d_loss = (d_fake_loss + d_real_loss) \n        d_loss = tf.reduce_sum(d_loss) / (batch_size * 2)\n\n    # get the gradients\n    gradients = d_tape.gradient(d_loss, model_D.trainable_variables)\n    \n    # update the weights of the discriminator\n    param_D.apply_gradients(zip(gradients, model_D.trainable_variables))\n    \n\n    # PHASE TWO - train the generator\n    with tf.GradientTape() as g_tape:\n        # create noise input\n        z = tf.random.normal(shape=(real_img.shape[0], 1, 1, z_dim))\n        \n        # generate fake images\n        fake_img = model_G(z)\n\n        # feed fake images to the discriminator\n        fake_out = model_D(fake_img)\n        \n        # use loss function to measure how well the generator\n        # is able to trick the discriminator (i.e. model_D should output 1's)\n        g_loss = loss_func(tf.ones_like(fake_out), fake_out)\n        g_loss = tf.reduce_sum(g_loss) / (batch_size * 2)\n    \n    # get the gradients\n    gradients = g_tape.gradient(g_loss, model_G.trainable_variables)\n\n    # update the weights of the generator\n    param_G.apply_gradients(zip(gradients, model_G.trainable_variables))\n    \n    # return the losses and fake images for monitoring\n    return d_loss, g_loss, fake_img ","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:36:07.310205Z","iopub.execute_input":"2023-03-09T22:36:07.311229Z","iopub.status.idle":"2023-03-09T22:36:07.328743Z","shell.execute_reply.started":"2023-03-09T22:36:07.311199Z","shell.execute_reply":"2023-03-09T22:36:07.327321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The Training loop.\n","metadata":{}},{"cell_type":"code","source":"NUM_EPOCHS = 100\n\n# generate a batch of noisy input\nz_dim = 128\ntest_z = tf.random.normal(shape=(64, 1, 1, z_dim))\n\n# start loop\ntf.keras.backend.clear_session()\nfor epoch in range(NUM_EPOCHS): \n    with tqdm(dataset) as pbar:\n        pbar.set_description(f\"[Epoch {epoch}]\")\n        for step, (X1, X2) in enumerate(pbar):\n            # train on the current batch\n            d_loss, g_loss, fake = train_on_batch(X1, X2)\n    \n        # generate fake images\n        fake_img = model_G(test_z)\n\n    # save face generated to file. \n    if not os.path.exists(out_dir):\n        os.makedirs(out_dir)\n    file_path = out_dir+f\"/epoch_{epoch:04}.png\"\n    \n    # display gallery of generated faces\n    if epoch % 1 == 0:\n        plot_results(fake_img.numpy()[:4], 2, save_path=file_path)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-03-09T22:36:07.330369Z","iopub.execute_input":"2023-03-09T22:36:07.3307Z","iopub.status.idle":"2023-03-09T22:40:44.157486Z","shell.execute_reply.started":"2023-03-09T22:36:07.330668Z","shell.execute_reply":"2023-03-09T22:40:44.156565Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image\nimgs = os.listdir('celeba_out')\nimgs.sort()\ngif_images = []\nfor img in imgs: \n  if img.find(\"png\")>0:\n    #print (\"celebaout/\" + img)\n    \n    with Image.open(\"celeba_out/\" + img) as im:\n      im = im.convert(\"RGB\")\n      gif_images.append(np.array(im))\ngif_path = \"faces.gif\"\nimageio.mimsave(gif_path, gif_images, fps=1)","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:40:44.158749Z","iopub.execute_input":"2023-03-09T22:40:44.159373Z","iopub.status.idle":"2023-03-09T22:41:04.173951Z","shell.execute_reply.started":"2023-03-09T22:40:44.159333Z","shell.execute_reply":"2023-03-09T22:41:04.171881Z"},"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"path=\"faces.gif\"\n\nwith open(path,'rb') as f:\n    display(IpyImage(data=f.read(), format='png'))","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:41:04.175372Z","iopub.execute_input":"2023-03-09T22:41:04.176033Z","iopub.status.idle":"2023-03-09T22:41:04.502562Z","shell.execute_reply.started":"2023-03-09T22:41:04.175979Z","shell.execute_reply":"2023-03-09T22:41:04.501236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from IPython.display import Image\ngif_path = \"faces.gif\"\nImage(filename=gif_path)","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:41:04.5038Z","iopub.execute_input":"2023-03-09T22:41:04.504102Z","iopub.status.idle":"2023-03-09T22:41:04.826511Z","shell.execute_reply.started":"2023-03-09T22:41:04.504076Z","shell.execute_reply":"2023-03-09T22:41:04.825713Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Conclusions, Fork & Improve. \n## If you liked the notebook, please consider **upvoting it**. I will be very happy, and it encourages me to continue publishing.\n\nPlease don't hesitate to fork, improve and play with the notebook at your convenience.\n\nThe faces you are looking at have been generated using only 1,000 images from the Dataset. If you use the full Dataset, the quality, and variety of faces will improve. \n\nWe've seen how to use TPUs and distributed execution in a notebook. I used it to train a GAN, but you can use it for any problem you want. \n\nThe main steps are: \n\n* Create a  strategy. \n* Instantiate the TPU. \n* Use the Scope to create the models. \n* Take care of the variables returned by functions executed in multiple TPUs. \n\nIf you like, try to convert the notebook to run on a GPU without distributed execution, and check the differences in performance. \n\nI hope you liked the notebook! \n\nOther GANs tutorials: \n\n**First DCGAN using TensorFlow:**\nhttps://www.kaggle.com/code/peremartramanonellas/gan-tutorial-first-dcgan-using-tensorflow\n\n**Generating Color Images:**\nhttps://www.kaggle.com/code/peremartramanonellas/gan-tutorial-2-generating-color-images\n\n**","metadata":{"execution":{"iopub.status.busy":"2023-03-09T22:13:55.738703Z","iopub.execute_input":"2023-03-09T22:13:55.739113Z","iopub.status.idle":"2023-03-09T22:13:55.74279Z","shell.execute_reply.started":"2023-03-09T22:13:55.739073Z","shell.execute_reply":"2023-03-09T22:13:55.742062Z"}}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}