{"cells":[{"metadata":{},"cell_type":"markdown","source":"# ***Disclaimer:*** \nHello Kagglers! I am a Solution Architect with the Google Cloud Platform. I am not an insider to this competition, so I am allowed to contribute and even compete, although I cannot collect prizes. The focus of my contributions is on helping users to leverage GCP components (GCS, TPUs, BigQueryetc..) in order to solve large problems. My ideas and contributions represent my own opinion, and are not representative of an official recommendation by Google. Also, I try to develop notebooks quickly in order to help users early in competitions. There may be better ways to solving particular problems, I welcome comments and suggestions. Use my contributions at your own risk, I don't garantee that they will help on winning any competition, but I am hoping to learn by collaborating with everyone.\n\n"},{"metadata":{},"cell_type":"markdown","source":"# Note: \nThis is a very early version of this Notebook. I got to the point that it uses distributed training on a very large dataset using TPUs, so I wanted to share it early to help folks dealing with the huge size of the RSNA-STR competition dataset. Check later versions for more comments and explanations, in particular are this point I am not checkpoing nor saving the trained model.\n\nThis notebook uses a sample TFRecords dataset which I made public so that everyone can try this example quickly. I am also publishing the notebook I used to create the dataset, which you can quickly modify to suit your needs. Here are the references for the input dataset and the Notebooks that explain how it was built:\n\nInput Dataset: [RSNA_PE_WINDOW_MIXED](https://www.kaggle.com/marcosnovaes/rsna-pe-window-mixed): Contains 10000 images from the competition dataset, 50% PE positive and 50% PE negative, labeled 1 for positive PE and 0 for negative. \n\nNotebook: [Building a TF Record Dataset](https://www.kaggle.com/marcosnovaes/building-a-tfrecord-dataset) : Explains how the TFRecord dataset was built\n\nNotebook: [Managing Large Datasets with TFRecords](https://www.kaggle.com/marcosnovaes/managing-large-datasets-with-tfrecords-on-gcp) : Explains in detail the how to build datasets with TFRecords"},{"metadata":{},"cell_type":"markdown","source":"**Objective**\nThis notebook explains how to deal with large datasets and also how to leverage distributed training with TPUs. This notebook can run using either GPUs or TPUs. If you are using GPUs you will have to limit the size of the dataset of you will ran out of memory. However, with TPUs the Notebook can process all 10000 images at full resolution to train a deep network. \n\n**Setup**\n1. This Notebook uses both Google Cloud Services and the Google Cloud SDK. Make sure to enable both by linking your GCP project using the Menu \"Add-ons-->Google Cloud SDK\" and \"Add-ons-->Google Cloud Services\".\n2. This Notebook must be configured with an accelerator. Add either a GPU or a TPU (works much better on TPUs)\n3. Add the sample TF Record dataset RSNA_PE_WINDOW_MIXED to the Notebook by using the \"Add\" button in the input folder of this Notebook"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd \nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nimport scipy.ndimage\n\nfrom os import listdir, mkdir\nimport os\nimport time\n\nimport tensorflow as tf\n\nfrom kaggle_datasets import KaggleDatasets\nfrom kaggle_secrets import UserSecretsClient\n\n# Import Keras Libraries\n\nfrom keras.models import Sequential\nfrom keras.layers import Dense\nfrom keras.layers import Reshape\nfrom keras.layers import Flatten\nfrom keras.layers import Conv2D\nfrom keras.layers import Conv2DTranspose\nfrom keras.layers import LeakyReLU\nfrom keras.layers import Dropout\nfrom tensorflow.keras import layers\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# If you are using TPUs, execute this cell and skip the next cell\n\n# Use the cluster resolver to communicate with the TPU\n\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\n\nstrategy = tf.distribute.experimental.TPUStrategy(tpu)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"worker_list=tpu.cluster_spec()\n#print(vars(worker_list))\nprint(worker_list._cluster_spec[\"worker\"][0])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!echo $TPU_NAME","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# If you are using a GPU, remove the comment and execute this cell as opposed to the previous cell\n#strategy = tf.distribute.MirroredStrategy()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Check how many cores we have. For TPUs we should see 8 cores, but only one for GPUs. The data will be distributed to all the cores, so we will use the number of cores to calculate the batch size for each instance"},{"metadata":{"trusted":true},"cell_type":"code","source":"strategy.num_replicas_in_sync","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"You must provide the name of a GCP project in this cell. This is just used to set up the storage API, which will be used to list the files that are stored in the dataset. YOU MUST CHANGE THE \"YOUR_PROJECT_ID\" TO POINT TO THE PROJECT YOU LINKED TO THIS NOTEBOOK USING THE MENU Add Ons-->Google Cloud Services"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Set your own project id here\nYOUR_PROJECT_ID = 'your_project_ID_here'\n\nfrom google.cloud import bigquery\nbigquery_client = bigquery.Client(project=YOUR_PROJECT_ID)\nfrom google.cloud import storage\nstorage_client = storage.Client(project=YOUR_PROJECT_ID)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Utility functions to interact with the Storage API"},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_bucket(dataset_name):\n    \"\"\"Creates a new bucket. https://cloud.google.com/storage/docs/ \"\"\"\n    bucket = storage_client.create_bucket(dataset_name)\n    print('Bucket {} created'.format(bucket.name))\n\ndef upload_blob(bucket_name, source_file_name, destination_blob_name):\n    \"\"\"Uploads a file to the bucket. https://cloud.google.com/storage/docs/ \"\"\"\n    bucket = storage_client.get_bucket(bucket_name)\n    blob = bucket.blob(destination_blob_name)\n    blob.upload_from_filename(source_file_name)\n    print('File {} uploaded to {}.'.format(\n        source_file_name,\n        destination_blob_name))\n    \ndef list_blobs(bucket_name):\n    \"\"\"Lists all the blobs in the bucket. https://cloud.google.com/storage/docs/\"\"\"\n    blobs = storage_client.list_blobs(bucket_name)\n    for blob in blobs:\n        print(blob.name)\n        \ndef get_blob_names(bucket_name):\n    \"\"\"Lists all the blobs in the bucket. https://cloud.google.com/storage/docs/\"\"\"\n    blobs = storage_client.list_blobs(bucket_name)\n    return blobs\n        \ndef download_to_kaggle(bucket_name,destination_directory,file_name):\n    \"\"\"Takes the data from your GCS Bucket and puts it into the working directory of your Kaggle notebook\"\"\"\n    os.makedirs(destination_directory, exist_ok = True)\n    full_file_path = os.path.join(destination_directory, file_name)\n    blobs = storage_client.list_blobs(bucket_name)\n    for blob in blobs:\n        blob.download_to_filename(full_file_path)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Code to read and decode TFRecords. This is the same code explained in the previous notebooks (see note at the top of this Notebook). I have made one change for this Notebook, I am applying the CT Window function as part of the decoding function. I had to change the CT Window function to use the Tensorflow math package (clip, min and max) as opposed to numpy. Numpy does not operate with Tensors, so all operations that you execute on a TPU that do computation must use Tensorflow equivalents. \n\nTODO on this version: Note that I commented out decoding the string values of the dataset (study Id and Image name). TPUs do not support passing string types on the dataset. It is still helpful to pass an identifier in case that we need other image attributes during training. This will be shown in a future version."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define the TFExample Data type for training models\n# Our TFRecord format will include the CT Image and metadata of the image, including the prediction label (is PE present)\n\n# Utilities serialize data into a TFRecord\ndef _bytes_feature(value):\n  \"\"\"Returns a bytes_list from a string / byte.\"\"\"\n  if isinstance(value, type(tf.constant(0))):\n    value = value.numpy() # BytesList won't unpack a string from an EagerTensor.\n  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))\n\ndef _float_feature(value):\n  \"\"\"Returns a float_list from a float / double.\"\"\"\n  return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))\n\ndef _int64_feature(value):\n  \"\"\"Returns an int64_list from a bool / enum / int / uint.\"\"\"\n  return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))\n\n# Create a dictionary describing the features.\nimage_feature_description = {\n    'height': tf.io.FixedLenFeature([], tf.int64),\n    'width': tf.io.FixedLenFeature([], tf.int64),\n    'image_raw': tf.io.FixedLenFeature([], tf.string),\n    'study_id': tf.io.FixedLenFeature([], tf.string),\n    'img_name': tf.io.FixedLenFeature([], tf.string),\n    'pred_label': tf.io.FixedLenFeature([], tf.int64)\n}\n\nPE_WINDOW_LEVEL = 100\nPE_WINDOW_WIDTH = 700\n\ndef _parse_image_function(example_proto):\n  # Parse the input tf.Example proto using the dictionary above.\n    single_example = tf.io.parse_single_example(example_proto, image_feature_description)\n    img_height = single_example['height']\n    img_width = single_example['width']\n    img_bytes = tf.io.decode_raw(single_example['image_raw'],out_type='float64')\n    resized_image = tf.reshape(img_bytes, (img_height,img_width))\n    windowed_image = CT_window(resized_image, PE_WINDOW_LEVEL,PE_WINDOW_WIDTH )\n    sample_image = tf.reshape(windowed_image, (img_height,img_width,1))\n    mtd = dict()\n    mtd['width'] = single_example['width']\n    mtd['height'] = single_example['height']\n    #mtd['study_id'] = tf.io.decode_base64(single_example['study_id'])\n    #mtd['img_name'] = tf.io.decode_base64(single_example['img_name'])\n    mtd['pred_label'] = single_example['pred_label']\n    struct = {\n    'img': sample_image,\n    'img_mtd': mtd\n    } \n    return struct\n\n\ndef read_tf_dataset(storage_file_path):\n    encoded_image_dataset = tf.data.TFRecordDataset(storage_file_path, compression_type=\"GZIP\")\n    record_structs = encoded_image_dataset.map(_parse_image_function)\n    return record_structs\n\ndef CT_window(img, WL=50, WW=350):\n    upper, lower = WL+WW//2, WL-WW//2\n    X = tf.clip_by_value(img, lower, upper)\n    X = X - tf.math.reduce_min(X)\n    X = X / tf.math.reduce_max(X)\n    return X","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Important Globals:\n1. BUFFER_SIZE = Used for shuffling the dataset (see tf.dataset.buffer() )\n1. GLOBAL_BATCH_SIZE = how many images in a training step. We will distributed them among the cores.(see tf.dataset.batch())\n1. BATCH_SIZE_PER_REPLICA = how many images each core will process in one training step"},{"metadata":{"trusted":true},"cell_type":"code","source":"## GLOBALS AND CONSTANTS\nBUFFER_SIZE=120\nGLOBAL_BATCH_SIZE=1024\nNUM_REPLICAS=strategy.num_replicas_in_sync\nBATCH_SIZE_PER_REPLICA = GLOBAL_BATCH_SIZE // NUM_REPLICAS\n\nEPOCHS=3\nLEARNING_RATE=0.00005\nBETA_1=0.1\n\n# provide a file name where checkpoints will be stored.\nexperiment_number = '2'\n# setup a bucket for saving checkpoints\nyour_bucket_name = 'your_bucket_name_here'\ncheckpoint_path = 'gs://'+your_bucket_name+'/training_checkpoints/exp' + experiment_number\ncheckpoint_prefix  = checkpoint_path + '/_ckpt'\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Grant Tensorflow permission to read the dataset. I ended up making my dataset public, but you can also use your own private dataset as long as you grant permissions as below. "},{"metadata":{"trusted":true},"cell_type":"code","source":"from kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nuser_credential = user_secrets.get_gcloud_credential()\nuser_secrets.set_tensorflow_credential(user_credential)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Discover the GCS path of the public dataset 'rsna-pe-window-mixed'. To make this easier, import the dataset using the \"Add\" button in the input folder of the Notebook. Then look for your dataset spelled exactly as it appears in the input folder (note that it uses \"-\" as opposed to the \"_\" separator that I used when creating the dataset."},{"metadata":{"trusted":true},"cell_type":"code","source":"#from kaggle_datasets import KaggleDatasets\nGCS_PATH = KaggleDatasets().get_gcs_path('rsna-pe-window-mixed')\nGCS_PATH","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"To get the name of the bucket that the dataset reside, you strip the \"gs://\" prefix"},{"metadata":{"trusted":true},"cell_type":"code","source":"# get a list of directories\n# strip the first 5 chars, that is the \"gs://\" prefix\nbucket_name = GCS_PATH[5:]\nbucket_name\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We now  make an array of filenames using the list_blobs function of the storage API. We put the \"gs://\" prefix back in the file names as we will need it when we give the filenames to the TPU. I divided the 10,000 images into 200 files in 5 directories. The reason is so that each tfrecord file will be around 20M. Dividing your dataset into smaller files will allow the tf.dataset object to do a better job of accessing them and managing its cache."},{"metadata":{"trusted":true},"cell_type":"code","source":"filenames = []\nblobs = storage_client.list_blobs(bucket_name)\nfor blob in blobs:\n    filenames.append('gs://{}/{}'.format(bucket_name,blob.name))\nfilenames.__len__()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"filenames[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Read the dataset locally for verification"},{"metadata":{"trusted":true},"cell_type":"code","source":"g_dataset = read_tf_dataset(filenames)\nsubset = g_dataset.take(1)\ntest_image = []\nfor struct in subset.as_numpy_iterator():\n    #struct = g_dataset.get_next()\n    img_mtd = struct[\"img_mtd\"]\n    img_bytes = struct[\"img\"]\n    test_image = img_bytes.reshape(512,512)\n    #print(\"img_name = {}, pred_label = {}, image_shape = {}\".format(img_mtd[\"img_name\"], img_mtd[\"pred_label\"], img_bytes.shape))\n    fig, ax = plt.subplots(1,2,figsize=(20,3))\n    ax[0].set_title(\"PE Specific CT-scan\")\n    ax[0].imshow(test_image, cmap=\"bone\")\n    ax[1].set_title(\"Pixelarray distribution\");\n    sns.distplot(test_image.flatten(), ax=ax[1]);\n    #for img_data in g_dataset[\"mtd\"]:\n    #    print(\"img_name = {}, pred_label = {}\".format(img_data[\"img_name\"], img_data[\"pred_label\"]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here is the classification model. This example is adapted from the [DC-GAN Tutorial](https://www.tensorflow.org/tutorials/generative/dcgan). It downsamples the image by 2 (stride = 2) using Conv2D. I am just using this model as an example, the focus of this Notebook is not model architecture. However, I find this model converges very quickly, so it lookd like a good choice. "},{"metadata":{"trusted":true},"cell_type":"code","source":"## Define the Discriminator Network - This example is for a 512 x 512 image. \n\ndef make_discriminator_model():\n    model = tf.keras.Sequential()\n    #model.add(layers.Conv2D(64, (5, 5), strides=(2), padding='same',input_shape=[512, 512, 1]))\n    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[512, 512, 1]))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n\n    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    model.add(layers.Conv2D(256, (5, 5), strides=(2, 2), padding='same'))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    model.add(layers.Conv2D(512, (5, 5), strides=(2,2 ), padding='same'))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n\n    model.add(layers.Flatten())\n    #model.add(layers.Dense(1))\n    model.add(layers.Dense(1, activation='relu'))\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The test image we loaded is of shape 512x512. To feed it into the model we need to add a batch dimension at position [0] and a channel dimension at position [3]"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_image.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_img = test_image.reshape(1,512,512,1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now test if the model works locally"},{"metadata":{"trusted":true},"cell_type":"code","source":"discriminator = make_discriminator_model()\n# provide the image we just generated, and get the decisio score. It should be near zero, since we provided a noie image\ndecision = discriminator(test_img)\nprint (decision)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Starting here, we will define some GLOBAL DISTRIBUTED variables and functions. By using the strategy.scope() context we make these functions and variables available to the TPUs. We will implement a [custom training loop as described in this tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training). "},{"metadata":{"trusted":true},"cell_type":"code","source":"# define loss function\nwith strategy.scope():\n    loss_object = tf.keras.losses.BinaryCrossentropy(\n    from_logits=True,\n    reduction=tf.keras.losses.Reduction.NONE)\n\n    def compute_loss(labels, predictions):\n        per_example_loss = loss_object(labels, predictions)\n        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define loss metrics\n\nwith strategy.scope():\n    test_loss = tf.keras.metrics.Mean(name='test_loss')\n\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Create a GLOBAL variable for the model. Note that the \"discriminator\" variable will be common for BOTH the CPU and TPU. This is how we can retrieve and save the trained model. Tensorflow takes care of all the serialization needed to make this happen, it is really an AWESOME feature."},{"metadata":{"trusted":true},"cell_type":"code","source":"# create global variables in the strategy scope\nwith strategy.scope():\n    discriminator = make_discriminator_model()\n    discriminator_optimizer = tf.keras.optimizers.Adam(lr=LEARNING_RATE, beta_1=BETA_1, amsgrad=False)\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The \"train_step\" function is the main function of the training loop. Note the \"tf.function\" adornment. This will activate the \"Autograph\" feature of TF 2.X and it will compile the function for fast execution in an accelerator. Once it is on the accelerator it is hard to debug, you will not have access to the print function. When developing, I first comment out the \"@tf.function\" to run locally, then I use @tf.function but still run locally on GPUs to test the compilation, then I finally switch to TPUs when everything works fine. "},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n        \n    def restore_models( gan_checkpoint, checkpoint_directory ):\n        status = gan_checkpoint.restore(tf.train.latest_checkpoint(checkpoint_directory))\n        return status\n\n# Define Train Step\n    @tf.function\n    def train_step(tf_records):\n        with tf.GradientTape() as disc_tape:\n            images = tf_records[\"img\"]\n            reshaped_images = tf.reshape(images,(BATCH_SIZE_PER_REPLICA,512,512,1))\n            labels = tf_records[\"img_mtd\"][\"pred_label\"]\n            labels = tf.dtypes.cast(labels, tf.float32)\n            labels = tf.reshape(labels, (BATCH_SIZE_PER_REPLICA,1))\n            model_output = discriminator(reshaped_images, training=True)\n            loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True,reduction=tf.keras.losses.Reduction.NONE)\n            per_example_loss = loss_object( labels, model_output)\n            disc_loss = tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)\n           \n        step_loss = 0.1 * disc_loss  \n        gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n        discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))\n       \n        return step_loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Note how these variable are used. "},{"metadata":{"trusted":true},"cell_type":"code","source":"GLOBAL_BATCH_SIZE","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH_SIZE_PER_REPLICA","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here is the call that execute the train step on all 8 cores of the TPU. This is done using strategy.run() followed by strategy.reduce() that will use the new \"All Reduce\" capability that was introduced for TPUs with TF version 2.X"},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n  # `run` replicates the provided computation and runs it\n  # with the distributed input. Note the use of the ReduceOp.SUM in the strategy.reduce operation\n    @tf.function\n    def distributed_train_step(dataset_inputs):\n        per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))\n        #test_loss = 0.5\n        reduced_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,axis=None)\n        return reduced_loss\n    \n    #tf.function\n    def distributed_test_step(dataset_inputs):\n        return strategy.run(test_step, args=(dataset_inputs,))\n\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The train loop is a local function that will iterate for each epoch. A Epoch is a run through the entire dataset. For testing purposes I am taking just 32 records as the \"subset\" variable, but then you have to set GLOBAL_BATCH_SIZE to a smaller number. By default the Notebook is using a GLOBAL_BATCH_SIZE = 1024, which means 128 images per step. If you are using GPU, you have to go MUCH smaller or you will run out of memory."},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop( num_epochs, input_dataset ):\n        for epoch in range(num_epochs):\n            total_loss = 0.0\n            num_batches = 0.0\n            start_time = time.time()\n           \n            # if you are testing, use a subset\n            #subset = input_dataset.take(32)\n            #train_dataset = subset.batch(GLOBAL_BATCH_SIZE, drop_remainder=True).cache()\n            \n            # when you are ready to run an epoch on the whole 10,000 images comment the line below. \n            train_dataset = input_dataset.batch(GLOBAL_BATCH_SIZE, drop_remainder=True).cache()\n            \n            # Distribute the dataset among the several cores\n            dist_train_dataset = strategy.experimental_distribute_dataset(train_dataset)\n            for x in dist_train_dataset:\n                total_loss += distributed_train_step(x)\n                # print a * for every step of GLOBAL_BATH_SIZE images\n                print(\"*\",end='')\n                num_batches += 1\n            train_loss = total_loss / num_batches\n            end_time = time.time()\n            elapsed_time = end_time - start_time \n            img_per_second = num_batches * GLOBAL_BATCH_SIZE /elapsed_time\n            # print how many images/sec we processed\n            print(\"training speed = {} images per second\".format(img_per_second))\n            \n        return elapsed_time, train_loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The \"train\" function is basically doing some house keeping for every iteration. It will periodically print progress every \"status_interval\" epoch and checkpoint (not implemented yet) after num_epochs."},{"metadata":{"trusted":true},"cell_type":"code","source":"# local function to drive training\ndef train( num_epochs, status_interval, check_option):\n\n    num_processed = 0\n    input_dataset = read_tf_dataset(filenames)\n    while num_processed < num_epochs:\n        print(\"Training Epoch #{}\".format(num_processed))\n        epoch_time, train_loss = train_loop(status_interval, input_dataset)\n        template = (\"Epoch {}, Loss: {}, elapsed time in epoch: {}\")\n        print (template.format(num_processed, train_loss, epoch_time))\n        num_processed += status_interval\n        \n        if check_option == 1:\n            template = (\"saving checkpoint: Epoch {}, Loss: {}, elapsed time in last epoch: {}\")\n            print (template.format(num_processed, train_loss, epoch_time))\n            gan_checkpoint.save(gan_checkpoint_prefix)    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The line below will train 3 Epochs, printing an update at every iteration (second arg = 1) and not checkpointing (third arg = 0). I am still implementing checkpointing. For the 10,000 images an epoch is about 128 seconds, not bad for full 512x512 images.\n\nNOTE: the first time you execute the step function it will take some time to compile the needed functions. This is a one time, the following steps will run faster as the code is pre-compiled. "},{"metadata":{"trusted":true},"cell_type":"code","source":"# !train(2,1,0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip3 install --upgrade \"tensorflow==2.3.0\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip3 install --upgrade \"cloud-tpu-profiler==2.3.0\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!/opt/conda/bin/capture_tpu_profile --service_addr 10.0.0.2:8470 --logdir '/kaggle/working/' --duration_ms 50 --num_tracing_attempts 1  --verbosity -1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!which capture_tpu_profile","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import multiprocessing \nimport subprocess\n\nimport os \n  \ndef worker1(): \n    # printing process id \n    print(\"ID of process running worker1: {}\".format(os.getpid())) \n    command = '/opt/conda/bin/capture_tpu_profile'\n    flags = '--service_addr 10.0.0.2:8470 --logdir /kaggle/working/ --duration_ms 5000'\n    #os.system(\"/opt/conda/bin/capture_tpu_profile/capture_tpu_profile --service_addr 10.0.0.2:8470 --logdir '/kaggle/working/' --duration_ms 5000 > /kaggle/working/cmd_output\")\n    with open('/kaggle/working/output.txt', 'w') as f:\n        process = subprocess.run([command, '--service_addr', '10.0.0.2:8470', '--logdir', '/kaggle/working/','--duration_ms','500', '--num_tracing_attempts', '10'], stdout=f)\n        \n    \ndef worker2(): \n    # printing process id \n    print(\"ID of process running worker2: {}\".format(os.getpid()))\n    train(3,1,0)\n  \nif __name__ == \"__main__\": \n    # printing main program process id \n    print(\"ID of main process: {}\".format(os.getpid())) \n  \n    # creating processes \n    p1 = multiprocessing.Process(target=worker1) \n    #p2 = multiprocessing.Process(target=train, args=(1,1,0,)) \n  \n    # starting processes \n    p1.start() \n    #p2.start()\n    train(1,1,0)\n\n    # process IDs \n    print(\"ID of process p1: {}\".format(p1.pid)) \n    #print(\"ID of process p2: {}\".format(p2.pid)) \n    \n    # wait until processes are finished \n    p1.join() \n    #p2.join() \n    \n    # both processes finished \n    print(\"Both processes finished execution!\") \n  \n    # check if processes are alive \n    print(\"Process p1 is alive: {}\".format(p1.is_alive())) \n    #print(\"Process p2 is alive: {}\".format(p2.is_alive()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n!cat  /kaggle/working/output.txt","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I hope this works for everyone, let me know otherwise! Cheers!"}],"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}