{"cells":[{"metadata":{},"cell_type":"markdown","source":"# ***Disclaimer:*** \nHello Kagglers! I am a Solution Architect with the Google Cloud Platform. I am a coach for this competition, 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."},{"metadata":{},"cell_type":"markdown","source":"# Objective:\nThe objective of this notebook is to demonstrate how to feed a Dataset to a TPU Accelerator. In previous notebooks (see list below) I have build a TFRecord dataset using 1 512x512 tile per file. This resulted in thousands of files which severely hurt the TPU performance. It turns out that the recommended TFRecord file size for feeding TPUs is about 100M. So, I have re-packated the previously built TFRecord Dataset packing 256 tiles per TFRecord file using the following notebook:\n\n[https://www.kaggle.com/marcosnovaes/hubmap-building-datasets-for-tpus](https://www.kaggle.com/marcosnovaes/hubmap-building-datasets-for-tpus/)\n\nThe resulting dataset was made public so you can try this notebook by adding \"hubmap-large-records\" to this notebook.\n\nLink to input dataset: [https://www.kaggle.com/marcosnovaes/hubmap-large-records](https://www.kaggle.com/marcosnovaes/hubmap-large-records)\n\nI built 3 datasets with tiles of a specific type. The datasets prefixes indicate the types:\n\n1) \"gloms\": tiles that have gloms in them (the mask has >0 density)\n\n2) \"nogloms\": tiles that have no glos (the mask has 0 density)\n\n3) \"cropped\": all non white tiles (with and without gloms. Basically the sum of 1 and 2)\n\nThe Unet Keras model utilized is the one proposed by a [popular paper in biomedical image segementation](https://arxiv.org/abs/1505.04597), by (Olaf Ronneberger, Philipp Fischer, Thomas Brox).\n\nThe particular implementation used is the one proposed by  [Dr. Bradley Erickson](https://github.com/slowvak), available in the:  [The Magician's Corner repository](https://github.com/RSNA/MagiciansCorner/blob/master/UNetWithTensorflow.ipynb).\n\nThe basic modification that I have made to the implementation provided by Dr. Erickson is to enable the Tensorflow distributed training strategy (tf.strategy). You will notice that the function model.fit() is used within a strategy.scope(), so that it leverages either GPU or TPU acceleration. \n\nIn previous notebooks, I demonstrated how to read the competition data and produce a TFRecord dataset tiling the images in 512x512 tiles. \n\nPrevious Notebooks in this competition:\nhttps://www.kaggle.com/marcosnovaes/hubmap-unet-keras-model-fit-with-gpu: Similar Keras Unet model, but running from the local file system with GPUs\n\nhttps://www.kaggle.com/marcosnovaes/hubmap-3-unet-models-with-keras-cpu-gpu/: Investigates three implementations of the Unet model\n\nhttps://www.kaggle.com/marcosnovaes/hubmap-read-data-and-build-tfrecords/: Demonstrates how the TFRecord Dataset was built\n\nhttps://www.kaggle.com/marcosnovaes/hubmap-looking-at-tfrecords/: Explains how to read the data using the TFRecord Dataset"},{"metadata":{},"cell_type":"markdown","source":"# Setup\n1) Add the TFRecord Dataset as input to the notebook: Go to the Data section at the right, click \"add data\" and lof for the dataset: \"hubmap_train_test\"\n\n2) This Notebook also shows how to access a Kaggle dataset directly from Google Cloud Storage (GCS). To enable this feature, you need to link the Notebook to a GCS project, by going to the menu Add-ons-->Cloud SDK\n\n3) Add a TPU to the Notebooks before running. Go to the Settings panel on the right and add a TPU v3-8 Accelerator"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport sys\nimport glob\nimport random\nimport warnings\n\nimport numpy as np\nimport pandas as pd\n\nimport matplotlib.pyplot as plt\n\nfrom tqdm import tqdm\nfrom itertools import chain\nfrom skimage.io import imread, imshow, imread_collection, concatenate_images\nfrom skimage.transform import resize\nfrom skimage.morphology import label\n\nfrom tensorflow.keras.models import Model, load_model\nfrom tensorflow.keras.layers import Input\n\nfrom tensorflow.keras.layers import Conv2D, Conv2DTranspose\nfrom tensorflow.keras.layers import MaxPooling2D, UpSampling2D\nfrom tensorflow.keras.layers import concatenate\n\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint\nfrom tensorflow.keras import backend as K\n\nfrom tensorflow.keras import layers\n\nfrom keras.engine.topology import Layer\n\nfrom tensorflow.keras.optimizers import Adam\nfrom keras.utils.generic_utils import get_custom_objects\n\n\nfrom kaggle_datasets import KaggleDatasets\nfrom kaggle_secrets import UserSecretsClient\n\nimport tensorflow as tf","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ACCELERATOR_TYPE = 'TPU'\n#ACCELERATOR_TYPE = 'GPU'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if ACCELERATOR_TYPE == 'TPU':\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.MirroredStrategy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"user_secrets = UserSecretsClient()\nuser_credential = user_secrets.get_gcloud_credential()\nuser_secrets.set_tensorflow_credential(user_credential)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Create a dictionary describing the features.\nimage_feature_description = {\n    'height': tf.io.FixedLenFeature([], tf.int64),\n    'width': tf.io.FixedLenFeature([], tf.int64),\n    'num_channels': tf.io.FixedLenFeature([], tf.int64),\n    'img_bytes': tf.io.FixedLenFeature([], tf.string),\n    'mask': tf.io.FixedLenFeature([], tf.string),\n}\n\ndef _parse_image_and_masks_function(example_proto):\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    num_channels = single_example['num_channels']\n    \n    img_bytes =  tf.io.decode_raw(single_example['img_bytes'],out_type='uint8')\n    #dynamic shape\n    #img_array = tf.reshape( img_bytes, (img_height, img_width, num_channels))\n    #fixed shape\n    img_array = tf.reshape( img_bytes, (512, 512, 3))\n    \n    mask_bytes =  tf.io.decode_raw(single_example['mask'],out_type='bool')\n\n    mask = tf.reshape(mask_bytes, (512,512))\n    \n    #normalize images array and cast image and mask to float32\n    img_array = tf.cast(img_array, tf.float32) / 255.0\n    mask = tf.cast(mask, tf.float32)\n    return img_array, mask\n\ndef read_dataset(storage_file_path):\n    encoded_image_dataset = tf.data.TFRecordDataset(storage_file_path, compression_type=\"GZIP\")\n    parsed_image_dataset = encoded_image_dataset.map(_parse_image_and_masks_function)\n    return parsed_image_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## Create a dictionary describing the features.\nimage_feature_description2 = {\n    'img_bytes': tf.io.FixedLenFeature([], tf.string),\n    'mask': tf.io.FixedLenFeature([], tf.string),\n    'tile_y_anchor': tf.io.FixedLenFeature([], tf.int64),\n    'tile_x_anchor': tf.io.FixedLenFeature([], tf.int64),\n    'label': tf.io.FixedLenFeature([], tf.int64)\n}\n\ndef _parse_image_function2(example_proto, TILE_SIZE, NUM_CHANNELS):\n  # Parse the input tf.Example proto using the dictionary above.\n    single_example = tf.io.parse_single_example(example_proto, image_feature_description2)\n    \n    img_bytes =  tf.io.decode_raw(single_example['img_bytes'],out_type='uint8')\n   \n    img_array = tf.reshape( img_bytes, (TILE_SIZE, TILE_SIZE, NUM_CHANNELS))\n    img_array = tf.cast(img_array, tf.float32) / 255.0\n    \n    mask_bytes =  tf.io.decode_raw(single_example['mask'],out_type='bool')\n    \n    mask = tf.reshape(mask_bytes, (TILE_SIZE,TILE_SIZE))\n    mask = tf.cast(mask, tf.float32)\n    \n    mtd = dict()\n    mtd['tile_y_anchor'] = single_example['tile_y_anchor']\n    mtd['tile_x_anchor'] = single_example['tile_x_anchor']\n    mtd['label'] = single_example['label']\n   \n    return img_array, mask\n\ndef read_tf_dataset2(storage_file_path, TILE_SIZE, NUM_CHANNELS):\n    encoded_image_dataset = tf.data.TFRecordDataset(storage_file_path, compression_type=\"GZIP\")\n    #ds = ds.map(lambda x: fun(x, my_arg))\n    \n    #parsed_image_dataset = encoded_image_dataset.map(_parse_image_function(TILE_SIZE=TILE_SIZE, NUM_CHANNELS=NUM_CHANNELS))\n    parsed_image_dataset = encoded_image_dataset.map(lambda x: _parse_image_function2(x,TILE_SIZE=TILE_SIZE, NUM_CHANNELS=NUM_CHANNELS))\n    return parsed_image_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    def dice_coeff(y_true, y_pred):\n        # add epsilon to avoid a divide by 0 error in case a slice has no pixels set\n        # we only care about relative value, not absolute so this alteration doesn't matter\n        _epsilon = 10 ** -7\n        intersections = tf.reduce_sum(y_true * y_pred)\n        unions = tf.reduce_sum(y_true + y_pred)\n        dice_scores = (2.0 * intersections + _epsilon) / (unions + _epsilon)\n        return dice_scores\n\n    def dice_loss(y_true, y_pred):\n        loss = 1 - dice_coeff(y_true, y_pred)\n        return loss\n  \n    get_custom_objects().update({\"dice\": dice_loss})\n\n    class LayerNormalization (Layer) :\n    \n        def call(self, x, mask=None, training=None) :\n            axis = list (range (1, len (x.shape)))\n            x /= K.std (x, axis = axis, keepdims = True) + K.epsilon()\n            x -= K.mean (x, axis = axis, keepdims = True)\n            return x\n        \n        def compute_output_shape(self, input_shape):\n            return input_shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def magic_unet(act_fn = 'relu', init_fn = 'he_normal', width=512, height = 512, channels = 3): \n    inputs = Input((512,512,3))\n    act_fn = 'relu'\n    init_fn = 'he_normal'\n\n    # note we use linear function before layer normalization\n    conv1 = Conv2D(8, 5, activation = 'linear', padding = 'same', kernel_initializer = init_fn)(inputs)\n    conv1 = LayerNormalization()(conv1)\n    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)\n    conv2 = Conv2D(16, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(pool1)\n    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)\n    conv3 = Conv2D(32, 3, activation = 'linear', padding = 'same', kernel_initializer = init_fn)(pool2)\n    conv3 = LayerNormalization()(conv3)\n    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)\n    conv4 = Conv2D(64, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(pool3)\n    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)\n\n    conv5 = Conv2D(72, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(pool4)\n\n    up6 = Conv2D(64, 2, activation = 'linear', padding = 'same', kernel_initializer = init_fn)(UpSampling2D(size = (2,2))(conv5))\n    up6 = LayerNormalization()(up6)\n    merge6 = concatenate([conv4,up6], axis = 3)\n    conv6 = Conv2D(64, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(merge6)\n\n    up7 = Conv2D(32, 2, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(UpSampling2D(size = (2,2))(conv6))\n    merge7 = concatenate([conv3,up7], axis = 3)\n    conv7 = Conv2D(32, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(merge7)\n\n    up8 = Conv2D(16, 2, activation = 'linear', padding = 'same', kernel_initializer = init_fn)(UpSampling2D(size = (2,2))(conv7))\n    up8 = LayerNormalization()(up8)\n    merge8 = concatenate([conv2,up8], axis = 3)\n    conv8 = Conv2D(16, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(merge8)\n\n    up9 = Conv2D(8, 2, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(UpSampling2D(size = (2,2))(conv8))\n    merge9 = concatenate([conv1,up9], axis = 3)\n    conv9 = Conv2D(8, 3, activation = act_fn, padding = 'same', kernel_initializer = init_fn)(merge9)\n    conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)\n    model = Model(inputs = inputs, outputs = conv10)\n\n    return model\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls /kaggle/input","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls /kaggle/input/hubmap-reduced","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#GCS_PATH = KaggleDatasets().get_gcs_path('hubmap-large-records')\nGCS_PATH = KaggleDatasets().get_gcs_path('hubmap-reduced')\nGCS_PATH","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#local_train_gloms_files = glob.glob('/kaggle/input/hubmap-large-records/train_gloms*.*')\nlocal_train_gloms_files = glob.glob('/kaggle/input/hubmap-reduced/gloms*.*')\nlocal_train_no_gloms_files = glob.glob('/kaggle/input/hubmap-reduced/no_gloms*.*')\n#local_validation_gloms_files = glob.glob('/kaggle/input/hubmap-large-records/validation_gloms*.*')\n#local_train_cropped_files = glob.glob('/kaggle/input/hubmap-large-records/train_cropped*.*')\n#local_validation_cropped_files = glob.glob('/kaggle/input/hubmap-large-records/validation_cropped*.*')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls -l /kaggle/input/hubmap-reduced","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# build corresponding GCS file paths\ndef build_gcs_paths( bucket_path, local_file_paths):\n    gcs_paths = []\n    for file_path in local_file_paths:\n        file_name = file_path.split('/')\n        file_name = file_name[-1]\n        gcs_name = bucket_path+'/'+file_name\n        gcs_paths = np.append(gcs_paths,gcs_name)\n    return gcs_paths","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_files = []\ntrain_files = np.append(train_files, local_train_gloms_files)\ntrain_files = np.append(train_files, local_train_no_gloms_files)\ngcs_train_gloms_files = build_gcs_paths(GCS_PATH,train_files)\n\n#gcs_validation_gloms_files = build_gcs_paths(GCS_PATH,local_validation_gloms_files)\n#gcs_train_cropped_files = build_gcs_paths(GCS_PATH,local_train_cropped_files)\n#gcs_validation_cropped_files = build_gcs_paths(GCS_PATH,local_validation_cropped_files)\n\n#gcs_validation_cropped_files","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gcs_train_gloms_files","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if ACCELERATOR_TYPE == 'TPU':\n    #train_dataset = read_dataset(gcs_train_gloms_files)\n    train_dataset = read_tf_dataset2(gcs_train_gloms_files, 512,3)\n    #validation_dataset = read_dataset(gcs_validation_gloms_files)\n    #train_dataset = read_dataset(gcs_train_cropped_files)\n    #validation_dataset = read_dataset(gcs_validation_cropped_files)\nelse:\n    train_dataset = read_dataset(local_train_gloms_files)\n    validation_dataset = read_dataset(local_validation_gloms_files)\n\ntrain_image = []\ntrain_mask =[]\nfor image, mask in train_dataset.take(1):\n    train_image, train_mask = image, mask\n    \ntest_image = []\ntest_mask =[]\nfor image, mask in train_dataset.take(1):\n    test_image, test_mask = image, mask\n\nfig, ax = plt.subplots(2,2,figsize=(10,6))\nax[0][0].imshow(train_image)\nax[0][1].imshow(train_mask)\nax[1][0].imshow(test_image)\nax[1][1].imshow(test_mask)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():   \n    model = magic_unet()\n    model.compile(optimizer = Adam(lr = 1e-3), loss = 'dice', metrics=[dice_coeff])\n    \n    if ACCELERATOR_TYPE == 'TPU':\n        batch_size = 128\n        train_dataset = read_tf_dataset2(gcs_train_gloms_files, 512, 3)\n        train_dataset = train_dataset.shuffle(5000)\n        #train_dataset = read_dataset(gcs_train_cropped_files)\n        train_dataset = train_dataset.batch(batch_size, drop_remainder=True).cache().prefetch(2)\n        \n        #validation_dataset = read_dataset(gcs_validation_gloms_files)\n        #validation_dataset = read_dataset(gcs_validation_cropped_files)\n        #validation_dataset = validation_dataset.batch(8, drop_remainder=True).cache().prefetch(2)\n    else:\n        batch_size = 8\n        train_dataset = read_dataset(local_train_gloms_files)\n        train_dataset = train_dataset.batch(batch_size, drop_remainder=True).prefetch(2)\n        validation_dataset = read_dataset(local_validation_gloms_files)\n        validation_dataset = validation_dataset.batch(batch_size, drop_remainder=True).prefetch(2)\n    \n    #steps_per_epoch = 60000 // batch_size\n    steps_per_epoch = 100\n    validation_steps = 10000 // batch_size\n    \n    checkpointer = ModelCheckpoint('/kaggle/working/hubmap-keras-tpu.h5', verbose=1)\n    \n    if ACCELERATOR_TYPE == 'TPU':\n       # model.fit(train_dataset,batch_size=64, epochs=300, validation_data=validation_dataset,callbacks=[checkpointer])\n        model.fit(train_dataset,batch_size=64, epochs=300,callbacks=[checkpointer])\n    else:\n        model.fit(train_dataset,batch_size=1, epochs=10, validation_data=validation_dataset,callbacks=[checkpointer])\n    \n    model.save_weights(\"/kaggle/working/hubmap-tpu-cortex-200.h5\")\n          #steps_per_epoch=steps_per_epoch,\n          #validation_data=test_dataset, \n          #validation_steps=validation_steps)\n    #earlystopper = EarlyStopping(patience=5, verbose=1)\n    #checkpointer = ModelCheckpoint('/kaggle/working/model-hubmap.h5', verbose=1)\n\n    #results = unet_model.fit(train_dataset, batch_size=1, epochs=1, callbacks=[checkpointer])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls -l /kaggle/working","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_model = magic_unet()\npred_model.load_weights(\"/kaggle/working/hubmap-tpu-gloms-200.h5\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_image = []\ntest_mask = []\npred_mask = []\nfor image, mask in train_dataset.take(1):\n    test_image, test_mask = image, mask\n    pred_mask = pred_model.predict(test_image, verbose=1)\n    bool_mask = (pred_mask > 0.5)\n    \nfig, ax = plt.subplots(1,3,figsize=(20,3))\nax[0].imshow(test_image[0,:,:,:])\nax[1].imshow(test_mask[0,:,:])\nax[2].imshow(bool_mask[0,:,:,0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mask_density = np.count_nonzero(pred_mask)\nmask_density","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}