{"cells":[{"metadata":{},"cell_type":"markdown","source":"**This notebooks shows three ways of training a model on TPU:**\n1. Using Keras and model.fit()\n1. Using a custom training loop\n1. Using a custom training loop specifically optimized for TPU\n\n**Optimization that benefit all three models:**\n\n- use `dataset.batch(BATCH_SIZE, drop_remainder=True)`<br/>\n   The training dataset is infinitely repeated so drop_remainder=True should not be needed. However, whith the setting, Tensorflow produces batches of a known size and although XLA (the TPU compiler) can now handle variable batches, it is slightly faster on fixed batches.<br/>\n   On the validation dataset, this setting can drop some validation images. It is not the case here because the validation dataset happens to contain an integral number of batches.\n   \n**Optimizations specific to the TPU-optimized custom training loop:**\n\n- The training and validation step functions run multiple batches at once. This is achieved by placing a loop using `tf.range()` in the step function. The loop will be compiled to (thanks to `@tf.function`) and executed on TPU.\n- The validation dataset is made to repeat indefinitely because handling end-of-dataset exception in a TPU loop implemented with `tf.range()` is not yet possible. Validation is adjusted to always use exactly or more than the entire validation dataset. This could change numerics. It happens that in this example, the validation dataset is used exactly once per validation.\n- The validation dataset iterator is not reset between validation runs. Since the iterator is passed into the step function which is then compiled for TPU (thanks to `@tf.function`), passing a fresh iterator for every validation run would trigger a fresh recompilation. With a validation at the end of every epoch this would be slow.\n- Losses are reported through Keras metrics. It is possible to return values from step function and return losses in that way. However, in the optimized version of the custom training loop, using `tf.range()`, aggregating losses returned from multiple batches becomes impractical."},{"metadata":{},"cell_type":"markdown","source":"## my tests\n\n>ex1: base with ep18; EN B5; aug: rndflip\n* time: 46.3s loss: 0.0138 accuracy: 0.9972 val_loss: 0.1969 val_acc: 0.9585 lr: 3.68e-05 steps/val_steps: 99/29\n\n>ex2: ep:30\n* abort error after ep 21\n* time: 46.3s loss: 0.0119 accuracy: 0.9977 val_loss: 0.1880 val_acc: 0.9639 lr: 2.372e-05 steps/val_steps: 99/29\n\n>V1: deotte ep:30\n* dd\n\n>4: augmix ep:30\n* dd\n"},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"import math, re, os, time\nimport tensorflow as tf, tensorflow.keras.backend as K\nimport numpy as np\nfrom collections import namedtuple\nfrom matplotlib import pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nfrom sklearn.metrics import f1_score, precision_score, recall_score, confusion_matrix\nprint(\"Tensorflow version \" + tf.__version__)\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# TPU or GPU detection"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Competition data access\nTPUs read data directly from Google Cloud Storage (GCS). This Kaggle utility will copy the dataset to a GCS bucket co-located with the TPU. If you have multiple datasets attached to the notebook, you can pass the name of a specific dataset to the get_gcs_path function. The name of the dataset is the name of the directory it is mounted in. Use `!ls /kaggle/input/` to list attached datasets."},{"metadata":{"trusted":true},"cell_type":"code","source":"#GCS_DS_PATH = KaggleDatasets().get_gcs_path() # you can list the bucket with \"!gsutil ls $GCS_DS_PATH\"\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('flower-classification') # you can list the bucket with \"!gsutil ls $GCS_DS_PATH\"\nGCS_DS_PATH_EXT = KaggleDatasets().get_gcs_path('flower-classification-with-tpus-external-datas-v3')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Configuration"},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_SIZE = [224, 224] # At this size, a GPU will run out of memory. Use the TPU.\n                        # For GPU training, please select 224 x 224 px image size.\nEPOCHS = 30\nBATCH_SIZE = 16 * strategy.num_replicas_in_sync\n\nGCS_PATH_SELECT = { # available image sizes\n    192: GCS_DS_PATH + '/tfrecords-jpeg-192x192',\n    224: GCS_DS_PATH + '/tfrecords-jpeg-224x224',\n    331: GCS_DS_PATH + '/tfrecords-jpeg-331x331',\n    512: GCS_DS_PATH + '/tfrecords-jpeg-512x512'\n}\nGCS_PATH = GCS_PATH_SELECT[IMAGE_SIZE[0]]\n\nTRAINING_FILENAMES_IN17 = tf.io.gfile.glob(GCS_DS_PATH_EXT + '/*2017.rec')\nTRAINING_FILENAMES_IN18 = tf.io.gfile.glob(GCS_DS_PATH_EXT + '/*2018.rec')\nTRAINING_FILENAMES_IN19 = tf.io.gfile.glob(GCS_DS_PATH_EXT + '/*2019.rec')\nTRAINING_FILENAMES_TF = tf.io.gfile.glob(GCS_DS_PATH_EXT + '/tf_flowers.rec')\nTRAINING_FILENAMES_OX = tf.io.gfile.glob(GCS_DS_PATH_EXT + '/Oxford.rec')\n\nSKIP_VALIDATION = False\n\n#if not SKIP_VALIDATION:\n#    TRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec') + TRAINING_FILENAMES_IN17 + TRAINING_FILENAMES_IN18  + TRAINING_FILENAMES_IN19  + \\\n#        TRAINING_FILENAMES_TF # + TRAINING_FILENAMES_OX    \n#else:\n#    TRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec') + TRAINING_FILENAMES_IN17 + TRAINING_FILENAMES_IN18  + TRAINING_FILENAMES_IN19  + \\\n#        TRAINING_FILENAMES_TF  + tf.io.gfile.glob(GCS_PATH + '/val/*.tfrec') #+ TRAINING_FILENAMES_OX \n\nTRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec')\nVALIDATION_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/val/*.tfrec')\nTEST_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/test/*.tfrec') # predictions on this dataset should be submitted for the competition\n\nCLASSES = ['pink primrose',    'hard-leaved pocket orchid', 'canterbury bells', 'sweet pea',     'wild geranium',     'tiger lily',           'moon orchid',              'bird of paradise', 'monkshood',        'globe thistle',         # 00 - 09\n           'snapdragon',       \"colt's foot\",               'king protea',      'spear thistle', 'yellow iris',       'globe-flower',         'purple coneflower',        'peruvian lily',    'balloon flower',   'giant white arum lily', # 10 - 19\n           'fire lily',        'pincushion flower',         'fritillary',       'red ginger',    'grape hyacinth',    'corn poppy',           'prince of wales feathers', 'stemless gentian', 'artichoke',        'sweet william',         # 20 - 29\n           'carnation',        'garden phlox',              'love in the mist', 'cosmos',        'alpine sea holly',  'ruby-lipped cattleya', 'cape flower',              'great masterwort', 'siam tulip',       'lenten rose',           # 30 - 39\n           'barberton daisy',  'daffodil',                  'sword lily',       'poinsettia',    'bolero deep blue',  'wallflower',           'marigold',                 'buttercup',        'daisy',            'common dandelion',      # 40 - 49\n           'petunia',          'wild pansy',                'primula',          'sunflower',     'lilac hibiscus',    'bishop of llandaff',   'gaura',                    'geranium',         'orange dahlia',    'pink-yellow dahlia',    # 50 - 59\n           'cautleya spicata', 'japanese anemone',          'black-eyed susan', 'silverbush',    'californian poppy', 'osteospermum',         'spring crocus',            'iris',             'windflower',       'tree poppy',            # 60 - 69\n           'gazania',          'azalea',                    'water lily',       'rose',          'thorn apple',       'morning glory',        'passion flower',           'lotus',            'toad lily',        'anthurium',             # 70 - 79\n           'frangipani',       'clematis',                  'hibiscus',         'columbine',     'desert-rose',       'tree mallow',          'magnolia',                 'cyclamen ',        'watercress',       'canna lily',            # 80 - 89\n           'hippeastrum ',     'bee balm',                  'pink quill',       'foxglove',      'bougainvillea',     'camellia',             'mallow',                   'mexican petunia',  'bromelia',         'blanket flower',        # 90 - 99\n           'trumpet creeper',  'blackberry lily',           'common tulip',     'wild rose']                                                                                                                                               # 100 - 102\n\n# Learning rate schedule for TPU, GPU and CPU.\n# Using an LR ramp up because fine-tuning a pre-trained model.\n# Starting with a high LR would break the pre-trained weights.\n\nLR_START = 0.00001\nLR_MAX = 0.00005 * strategy.num_replicas_in_sync\nLR_MIN = 0.00001\nLR_RAMPUP_EPOCHS = 5\nLR_SUSTAIN_EPOCHS = 0\nLR_EXP_DECAY = .8\n        \n@tf.function\ndef lrfn(epoch):\n    if epoch < LR_RAMPUP_EPOCHS:\n        lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n    elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:\n        lr = LR_MAX\n    else:\n        lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n    return lr\n\nrng = [i for i in range(EPOCHS)]\ny = [lrfn(x) for x in rng]\nplt.plot(rng, y)\nprint(\"Learning rate schedule: {:.3g} to {:.3g} to {:.3g}\".format(y[0], max(y), y[-1]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualization utilities\ndata -> pixels, nothing of much interest for the machine learning practitioner in this section."},{"metadata":{"trusted":true},"cell_type":"code","source":"# numpy and matplotlib defaults\nnp.set_printoptions(threshold=15, linewidth=80)\n\ndef batch_to_numpy_images_and_labels(data):\n    images, labels = data\n    numpy_images = images.numpy()\n    numpy_labels = labels.numpy()\n    if numpy_labels.dtype == object: # binary string in this case, these are image ID strings\n        numpy_labels = [None for _ in enumerate(numpy_images)]\n    # If no labels, only image IDs, return None for labels (this is the case for test data)\n    return numpy_images, numpy_labels\n\ndef title_from_label_and_target(label, correct_label):\n    if correct_label is None:\n        return CLASSES[label], True\n    correct = (label == correct_label)\n    return \"{} [{}{}{}]\".format(CLASSES[label], 'OK' if correct else 'NO', u\"\\u2192\" if not correct else '',\n                                CLASSES[correct_label] if not correct else ''), correct\n\ndef display_one_flower(image, title, subplot, red=False, titlesize=16):\n    plt.subplot(*subplot)\n    plt.axis('off')\n    plt.imshow(image)\n    if len(title) > 0:\n        plt.title(title, fontsize=int(titlesize) if not red else int(titlesize/1.2), color='red' if red else 'black', fontdict={'verticalalignment':'center'}, pad=int(titlesize/1.5))\n    return (subplot[0], subplot[1], subplot[2]+1)\n    \ndef display_batch_of_images(databatch, predictions=None):\n    \"\"\"This will work with:\n    display_batch_of_images(images)\n    display_batch_of_images(images, predictions)\n    display_batch_of_images((images, labels))\n    display_batch_of_images((images, labels), predictions)\n    \"\"\"\n    # data\n    images, labels = batch_to_numpy_images_and_labels(databatch)\n    if labels is None:\n        labels = [None for _ in enumerate(images)]\n        \n    # auto-squaring: this will drop data that does not fit into square or square-ish rectangle\n    rows = int(math.sqrt(len(images)))\n    cols = len(images)//rows\n        \n    # size and spacing\n    FIGSIZE = 13.0\n    SPACING = 0.1\n    subplot=(rows,cols,1)\n    if rows < cols:\n        plt.figure(figsize=(FIGSIZE,FIGSIZE/cols*rows))\n    else:\n        plt.figure(figsize=(FIGSIZE/rows*cols,FIGSIZE))\n    \n    # display\n    for i, (image, label) in enumerate(zip(images[:rows*cols], labels[:rows*cols])):\n        title = '' if label is None else CLASSES[label]\n        correct = True\n        if predictions is not None:\n            title, correct = title_from_label_and_target(predictions[i], label)\n        dynamic_titlesize = FIGSIZE*SPACING/max(rows,cols)*40+3 # magic formula tested to work from 1x1 to 10x10 images\n        subplot = display_one_flower(image, title, subplot, not correct, titlesize=dynamic_titlesize)\n    \n    #layout\n    plt.tight_layout()\n    if label is None and predictions is None:\n        plt.subplots_adjust(wspace=0, hspace=0)\n    else:\n        plt.subplots_adjust(wspace=SPACING, hspace=SPACING)\n    plt.show()\n\ndef display_confusion_matrix(cmat, score, precision, recall):\n    plt.figure(figsize=(15,15))\n    ax = plt.gca()\n    ax.matshow(cmat, cmap='Reds')\n    ax.set_xticks(range(len(CLASSES)))\n    ax.set_xticklabels(CLASSES, fontdict={'fontsize': 7})\n    plt.setp(ax.get_xticklabels(), rotation=45, ha=\"left\", rotation_mode=\"anchor\")\n    ax.set_yticks(range(len(CLASSES)))\n    ax.set_yticklabels(CLASSES, fontdict={'fontsize': 7})\n    plt.setp(ax.get_yticklabels(), rotation=45, ha=\"right\", rotation_mode=\"anchor\")\n    titlestring = \"\"\n    if score is not None:\n        titlestring += 'f1 = {:.3f} '.format(score)\n    if precision is not None:\n        titlestring += '\\nprecision = {:.3f} '.format(precision)\n    if recall is not None:\n        titlestring += '\\nrecall = {:.3f} '.format(recall)\n    if len(titlestring) > 0:\n        ax.text(101, 1, titlestring, fontdict={'fontsize': 18, 'horizontalalignment':'right', 'verticalalignment':'top', 'color':'#804040'})\n    plt.show()\n    \ndef display_training_curves(training, validation, title, subplot):\n    if subplot%10==1: # set up the subplots on the first call\n        plt.subplots(figsize=(10,10), facecolor='#F0F0F0')\n        plt.tight_layout()\n    ax = plt.subplot(subplot)\n    ax.set_facecolor('#F8F8F8')\n    ax.plot(training)\n    ax.plot(validation)\n    ax.set_title('model '+ title)\n    ax.set_ylabel(title)\n    #ax.set_ylim(0.28,1.05)\n    ax.set_xlabel('epoch')\n    ax.legend(['train', 'valid.'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# AugMix"},{"metadata":{"trusted":true},"cell_type":"code","source":"def int_parameter(level, maxval):\n    return tf.cast(level * maxval / 10, tf.int32)\n\ndef float_parameter(level, maxval):\n    return tf.cast((level) * maxval / 10., tf.float32)\n\ndef sample_level(n):\n    return tf.random.uniform(shape=[1], minval=0.1, maxval=n, dtype=tf.float32)\n    \ndef affine_transform(image, transform_matrix):\n    # input image - is one image of size [dim,dim,3] not a batch of [b,dim,dim,3]\n    DIM = IMAGE_SIZE[0]\n    XDIM = DIM%2 #fix for size 331\n    \n    x = tf.repeat(tf.range(DIM//2,-DIM//2,-1), DIM)\n    y = tf.tile(tf.range(-DIM//2,DIM//2), [DIM])\n    z = tf.ones([DIM*DIM], dtype='int32')\n    idx = tf.stack([x, y, z])\n    \n    # ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS\n    idx2 = K.dot(transform_matrix, tf.cast(idx, dtype='float32'))\n    idx2 = K.cast(idx2, dtype='int32')\n    idx2 = K.clip(idx2, -DIM//2+XDIM+1, DIM//2)\n    \n    # FIND ORIGIN PIXEL VALUES           \n    idx3 = tf.stack([DIM//2-idx2[0,], DIM//2-1+idx2[1,]])\n    d = tf.gather_nd(image, tf.transpose(idx3))\n    return tf.reshape(d,[DIM,DIM,3])\n\ndef blend(image1, image2, factor):\n    if factor == 0.0:\n        return tf.convert_to_tensor(image1)\n    if factor == 1.0:\n        return tf.convert_to_tensor(image2)\n\n    image1 = tf.cast(image1, tf.float32)\n    image2 = tf.cast(image2, tf.float32)\n\n    difference = image2 - image1\n    scaled = factor * difference\n\n    # Do addition in float.\n    temp = tf.cast(image1, tf.float32) + scaled\n\n    # Interpolate\n    if factor > 0.0 and factor < 1.0:\n        # Interpolation means we always stay within 0 and 255.\n        return tf.cast(temp, tf.uint8)\n\n    # Extrapolate:\n    #\n    # We need to clip and then cast.\n    return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def rotate(image, level):\n    degrees = float_parameter(sample_level(level), 30)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    degrees = tf.cond(rand_var > 0.5, lambda: degrees, lambda: -degrees)\n\n    angle = math.pi*degrees/180 # convert degrees to radians\n    angle = tf.cast(angle, tf.float32)\n    # define rotation matrix\n    c1 = tf.math.cos(angle)\n    s1 = tf.math.sin(angle)\n    one = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n    rotation_matrix = tf.reshape(tf.concat([c1,s1,zero, -s1,c1,zero, zero,zero,one], axis=0), [3,3])\n\n    transformed = affine_transform(image, rotation_matrix)\n    return transformed\n\ndef translate_x(image, level):\n    lvl = int_parameter(sample_level(level), IMAGE_SIZE[0] / 3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    lvl = tf.cast(lvl, tf.float32)\n    translate_x_matrix = tf.reshape(tf.concat([one,zero,zero, zero,one,lvl, zero,zero,one], axis=0), [3,3])\n\n    transformed = affine_transform(image, translate_x_matrix)\n    return transformed\n\ndef translate_y(image, level):\n    lvl = int_parameter(sample_level(level), IMAGE_SIZE[0] / 3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    lvl = tf.cast(lvl, tf.float32)\n    translate_y_matrix = tf.reshape(tf.concat([one,zero,lvl, zero,one,zero, zero,zero,one], axis=0), [3,3])\n\n    transformed = affine_transform(image, translate_y_matrix)\n    return transformed\n\ndef shear_x(image, level):\n    lvl = float_parameter(sample_level(level), 0.3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    s2 = tf.math.sin(lvl)\n    shear_x_matrix = tf.reshape(tf.concat([one,s2,zero, zero,one,zero, zero,zero,one],axis=0), [3,3])   \n\n    transformed = affine_transform(image, shear_x_matrix)\n    return transformed\n\ndef shear_y(image, level):\n    lvl = float_parameter(sample_level(level), 0.3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    c2 = tf.math.cos(lvl)\n    shear_y_matrix = tf.reshape(tf.concat([one,zero,zero, zero,c2,zero, zero,zero,one],axis=0), [3,3])   \n    \n    transformed = affine_transform(image, shear_y_matrix)\n    return transformed\n\ndef solarize(image, level):\n    # For each pixel in the image, select the pixel\n    # if the value is less than the threshold.\n    # Otherwise, subtract 255 from the pixel.\n    threshold = float_parameter(sample_level(level), 1)\n    return tf.where(image < threshold, image, 1 - image)\n\ndef solarize_add(image, level):\n    # For each pixel in the image less than threshold\n    # we add 'addition' amount to it and then clip the\n    # pixel value to be between 0 and 255. The value\n    # of 'addition' is between -128 and 128.\n    threshold = float_parameter(sample_level(level), 1)\n    addition = float_parameter(sample_level(level), 0.5)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    addition = tf.cond(rand_var > 0.5, lambda: addition, lambda: -addition)\n\n    added_image = tf.cast(image, tf.float32) + addition\n    added_image = tf.cast(tf.clip_by_value(added_image, 0, 1), tf.float32)\n    return tf.where(image < threshold, added_image, image)\n\ndef posterize(image, level):\n    lvl = int_parameter(sample_level(level), 8)\n    shift = 8 - lvl\n    shift = tf.cast(shift, tf.uint8)\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n    image = tf.bitwise.left_shift(tf.bitwise.right_shift(image, shift), shift)\n    return tf.cast(tf.clip_by_value(tf.math.divide(image, 255), 0, 1), tf.float32)\n\ndef autocontrast(image, _):\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n\n    def scale_channel(image):\n        # A possibly cheaper version can be done using cumsum/unique_with_counts\n        # over the histogram values, rather than iterating over the entire image.\n        # to compute mins and maxes.\n        lo = tf.cast(tf.reduce_min(image), tf.float32)\n        hi = tf.cast(tf.reduce_max(image), tf.float32)\n\n        # Scale the image, making the lowest value 0 and the highest value 255.\n        def scale_values(im):\n            scale = 255.0 / (hi - lo)\n            offset = -lo * scale\n            im = tf.cast(im, tf.float32) * scale + offset\n            im = tf.clip_by_value(im, 0.0, 255.0)\n            return tf.cast(im, tf.uint8)\n\n        result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)\n        return result\n\n    # Assumes RGB for now.  Scales each channel independently\n    # and then stacks the result.\n    s1 = scale_channel(image[:, :, 0])\n    s2 = scale_channel(image[:, :, 1])\n    s3 = scale_channel(image[:, :, 2])\n    image = tf.stack([s1, s2, s3], 2)\n    return tf.cast(tf.clip_by_value(tf.math.divide(image, 255), 0, 1), tf.float32)\n\ndef equalize(image, _):\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n\n    def scale_channel(im, c):\n        im = tf.cast(im[:, :, c], tf.int32)\n        # Compute the histogram of the image channel.\n        histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)\n        # For the purposes of computing the step, filter out the nonzeros.\n        nonzero = tf.where(tf.not_equal(histo, 0))\n        nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])\n        step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255\n\n        def build_lut(histo, step):\n            # Compute the cumulative sum, shifting by step // 2\n            # and then normalization by step.\n            lut = (tf.cumsum(histo) + (step // 2)) // step\n            # Shift lut, prepending with 0.\n            lut = tf.concat([[0], lut[:-1]], 0)\n            # Clip the counts to be in range.  This is done\n            # in the C code for image.point.\n            return tf.clip_by_value(lut, 0, 255)\n\n        # If step is zero, return the original image.  Otherwise, build\n        # lut from the full histogram and step and then index from it.\n        result = tf.cond(tf.equal(step, 0),\n                        lambda: im,\n                        lambda: tf.gather(build_lut(histo, step), im))\n\n        return tf.cast(result, tf.uint8)\n\n    # Assumes RGB for now.  Scales each channel independently\n    # and then stacks the result.\n    s1 = scale_channel(image, 0)\n    s2 = scale_channel(image, 1)\n    s3 = scale_channel(image, 2)\n    image = tf.stack([s1, s2, s3], 2)\n\n    return tf.cast(tf.clip_by_value(tf.math.divide(image, 255), 0, 1), tf.float32)\n\ndef color(image, level):\n    factor = float_parameter(sample_level(level), 1.8) + 0.1\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n    degenerate = tf.image.grayscale_to_rgb(tf.image.rgb_to_grayscale(image))\n    blended = blend(degenerate, image, factor)\n    return tf.cast(tf.clip_by_value(tf.math.divide(blended, 255), 0, 1), tf.float32)\n\ndef brightness(image, level):\n    delta = float_parameter(sample_level(level), 0.5) + 0.1\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    delta = tf.cond(rand_var > 0.5, lambda: delta, lambda: -delta) \n    return tf.image.adjust_brightness(image, delta=delta)\n\ndef contrast(image, level):\n    factor = float_parameter(sample_level(level), 1.8) + 0.1\n    factor = tf.reshape(factor, [])\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    factor = tf.cond(rand_var > 0.5, lambda: factor, lambda: 1.9 - factor  )\n\n    return tf.image.adjust_contrast(image, factor)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"means = {'R': 0.44892993872313053, 'G': 0.4148519066242368, 'B': 0.301880284715257}\nstds = {'R': 0.24393544875614917, 'G': 0.2108791383467354, 'B': 0.220427056859487}\n\ndef substract_means(image):\n    image = image - np.array([means['R'], means['G'], means['B']])\n    return image\n\ndef normalize(image):\n    image = substract_means(image)\n    image = image / np.array([stds['R'], stds['G'], stds['B']])\n    return tf.clip_by_value(image, 0, 1)\n\ndef apply_op(image, level, which):\n    # is there any better way than manually typing all of these conditions? \n    # I tried to randomly select transformation from array of functions, but tensorflow didn't let me to\n    augmented = image\n    augmented = tf.cond(which == tf.constant([0], dtype=tf.int32), lambda: rotate(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([1], dtype=tf.int32), lambda: translate_x(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([2], dtype=tf.int32), lambda: translate_y(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([3], dtype=tf.int32), lambda: shear_x(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([4], dtype=tf.int32), lambda: shear_y(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([5], dtype=tf.int32), lambda: solarize_add(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([6], dtype=tf.int32), lambda: solarize(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([7], dtype=tf.int32), lambda: posterize(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([8], dtype=tf.int32), lambda: autocontrast(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([9], dtype=tf.int32), lambda: equalize(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([10], dtype=tf.int32), lambda: color(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([11], dtype=tf.int32), lambda: contrast(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([12], dtype=tf.int32), lambda: brightness(image, level), lambda: augmented)\n    return augmented\n\ndef augmix(image):\n    # you can play with these parameters\n    severity = 3 # level of transformations as described above in transformations (integer from 1 to 10)\n    width = 3 # number of different chains of transformations to be mixed\n    depth = -1 # number of transformations in one chain, -1 means random from 1 to 3\n    \n    alpha = 1.\n    dir_dist = tfp.distributions.Dirichlet([alpha]*width)\n    ws = tf.cast(dir_dist.sample(), tf.float32)\n    beta_dist = tfp.distributions.Beta(alpha, alpha)\n    m = tf.cast(beta_dist.sample(), tf.float32)\n\n    mix = tf.zeros_like(image, dtype='float32')\n\n    def outer_loop_cond(i, depth, mix):\n        return tf.less(i, width)\n\n    def outer_loop_body(i, depth, mix):\n        image_aug = tf.identity(image)\n        depth = tf.cond(tf.greater(depth, 0), lambda: depth, lambda: tf.random.uniform(shape=[], minval=1, maxval=3, dtype=tf.int32))\n\n        def inner_loop_cond(j, image_aug):\n            return tf.less(j, depth)\n\n        def inner_loop_body(j, image_aug):\n            which = tf.random.uniform(shape=[], minval=0, maxval=3, dtype=tf.int32)\n            image_aug = apply_op(image_aug, severity, which)\n            j = tf.add(j, 1)\n            return j, image_aug\n        \n        j = tf.constant([0], dtype=tf.int32)\n        j, image_aug = tf.while_loop(inner_loop_cond, inner_loop_body, [j, image_aug])\n\n        wsi = tf.gather(ws, i)\n        mix = tf.add(mix, wsi*normalize(image_aug))\n        i = tf.add(i, 1)\n        return i, depth, mix\n\n    i = tf.constant([0], dtype=tf.int32)\n    i, depth, mix = tf.while_loop(outer_loop_cond, outer_loop_body, [i, depth, mix])\n    \n    mixed = tf.math.scalar_mul((1 - m), normalize(image)) + tf.math.scalar_mul(m, mix)\n    return tf.clip_by_value(mixed, 0, 1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Deotte's spatial transforms"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_mat(rotation, shear, height_zoom, width_zoom, height_shift, width_shift):\n    # returns 3x3 transformmatrix which transforms indicies\n        \n    # CONVERT DEGREES TO RADIANS\n    rotation = math.pi * rotation / 180.\n    shear = math.pi * shear / 180.\n    \n    # ROTATION MATRIX\n    c1 = tf.math.cos(rotation)\n    s1 = tf.math.sin(rotation)\n    one = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n    rotation_matrix = tf.reshape( tf.concat([c1,s1,zero, -s1,c1,zero, zero,zero,one],axis=0),[3,3] )\n        \n    # SHEAR MATRIX\n    c2 = tf.math.cos(shear)\n    s2 = tf.math.sin(shear)\n    shear_matrix = tf.reshape( tf.concat([one,s2,zero, zero,c2,zero, zero,zero,one],axis=0),[3,3] )    \n    \n    # ZOOM MATRIX\n    zoom_matrix = tf.reshape( tf.concat([one/height_zoom,zero,zero, zero,one/width_zoom,zero, zero,zero,one],axis=0),[3,3] )\n    \n    # SHIFT MATRIX\n    shift_matrix = tf.reshape( tf.concat([one,zero,height_shift, zero,one,width_shift, zero,zero,one],axis=0),[3,3] )\n    \n    return K.dot(K.dot(rotation_matrix, shear_matrix), K.dot(zoom_matrix, shift_matrix))\n\n\ndef transformDeotteSpacial(image,label):\n    # input image - is one image of size [dim,dim,3] not a batch of [b,dim,dim,3]\n    # output - image randomly rotated, sheared, zoomed, and shifted\n    \n    DIM = IMAGE_SIZE[0]\n    XDIM = DIM%2 #fix for size 331\n    \n    rot = 15. * tf.random.normal([1],dtype='float32')\n    shr = 5. * tf.random.normal([1],dtype='float32') \n    h_zoom = 1.0 + tf.random.normal([1],dtype='float32')/10.\n    w_zoom = 1.0 + tf.random.normal([1],dtype='float32')/10.\n    #w_zoom = h_zoom\n    h_shift = 16. * tf.random.normal([1],dtype='float32') \n    w_shift = 16. * tf.random.normal([1],dtype='float32') \n  \n    # GET TRANSFORMATION MATRIX\n    m = get_mat(rot,shr,h_zoom,w_zoom,h_shift,w_shift) \n\n    # LIST DESTINATION PIXEL INDICES\n    x = tf.repeat( tf.range(DIM//2,-DIM//2,-1), DIM )\n    y = tf.tile( tf.range(-DIM//2,DIM//2),[DIM] )\n    z = tf.ones([DIM*DIM],dtype='int32')\n    idx = tf.stack( [x,y,z] )\n    \n    # ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS\n    idx2 = K.dot(m,tf.cast(idx,dtype='float32'))\n    idx2 = K.cast(idx2,dtype='int32')\n    idx2 = K.clip(idx2,-DIM//2+XDIM+1,DIM//2)\n    \n    # FIND ORIGIN PIXEL VALUES           \n    idx3 = tf.stack( [DIM//2-idx2[0,], DIM//2-1+idx2[1,]] )\n    d = tf.gather_nd(image,tf.transpose(idx3))\n    #print('transform')\n        \n    return tf.reshape(d,[DIM,DIM,3]),label","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Datasets"},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"def decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n    \n    image = tf.image.resize(image, [224, 224])\n    \n    image = tf.reshape(image, [*IMAGE_SIZE, 3]) # explicit size needed for TPU\n    \n    #image = tf.image.resize_with_crop_or_pad(image, target_height, target_width)\n    #image = tf.image.resize(image, size=[*IMAGE_SIZE], preserve_aspect_ratio=False,antialias=False, name=None)\n    #image = augment_and_mix(image)\n    \n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label # returns a dataset of (image, label) pairs\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"id\": tf.io.FixedLenFeature([], tf.string),  # shape [] means single element\n        # class is missing, this competitions's challenge is to predict flower classes for the test dataset\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['id']\n    return image, idnum # returns a dataset of image(s)\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    # Read from TFRecords. For optimal performance, reading from multiple files at once and\n    # disregarding data order. Order does not matter since we will be shuffling the data anyway.\n\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls=AUTO)\n    # returns a dataset of (image, label) pairs if labeled=True or (image, id) pairs if labeled=False\n    return dataset\n\ndef data_augment(image, label):\n    # data augmentation. Thanks to the dataset.prefetch(AUTO) statement in the next function (below),\n    # this happens essentially for free on TPU. Data pipeline code is executed on the \"CPU\" part\n    # of the TPU while the TPU itself is computing gradients.\n    image = tf.image.random_flip_left_right(image)\n    #image = tf.image.random_saturation(image, 0, 2)\n    return image, label   \n\ndef get_training_dataset():\n    dataset = load_dataset(TRAINING_FILENAMES, labeled=True)\n    #dataset = load_dataset(TRAINING_FILENAMES_IN17, labeled=True)\n    dataset = dataset.map(data_augment, num_parallel_calls=AUTO)\n    \n    dataset = dataset.map(transformDeotteSpacial, num_parallel_calls=AUTO)\n    \n    dataset = dataset.repeat() # the training dataset must repeat for several epochs\n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True) # slighly faster with fixed tensor sizes\n    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)\n    return dataset\n\ndef get_validation_dataset(ordered=False, repeated=False):\n    dataset = load_dataset(VALIDATION_FILENAMES, labeled=True, ordered=ordered)\n    if repeated:\n        dataset = dataset.repeat()\n        dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=repeated) # slighly faster with fixed tensor sizes\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)\n    return dataset\n\ndef get_test_dataset(ordered=False):\n    dataset = load_dataset(TEST_FILENAMES, labeled=False, ordered=ordered)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)\n    return dataset\n\n#def count_data_items(filenames):\n#    # the number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items\n#    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n#    return np.sum(n)\n\ndef resize(image,label):\n    image = tf.image.resize(image, [224, 224])\n    return image,label\n\ndef count_data_items(filenames, labeled=False):\n    dataset = load_dataset(filenames,labeled = labeled)\n    #dataset = dataset.map(resize, num_parallel_calls=AUTO)\n    counter = 0\n    for element in dataset.as_numpy_iterator(): \n        counter = counter +1\n        \n    return counter\n\ndef int_div_round_up(a, b):\n    return (a + b - 1) // b\n\nNUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES, labeled=True)\nNUM_VALIDATION_IMAGES = count_data_items(VALIDATION_FILENAMES, labeled=True)\nNUM_TEST_IMAGES = count_data_items(TEST_FILENAMES)\nSTEPS_PER_EPOCH = NUM_TRAINING_IMAGES // BATCH_SIZE\nVALIDATION_STEPS = int_div_round_up(NUM_VALIDATION_IMAGES, BATCH_SIZE)\nprint('Dataset: {} training images, {} validation images, {} unlabeled test images'.format(NUM_TRAINING_IMAGES, NUM_VALIDATION_IMAGES, NUM_TEST_IMAGES))\n\nNUM_TRAINING_IMAGES_EXT = count_data_items(TRAINING_FILENAMES_IN17, labeled=True)\nprint('IN17; {} images'.format(NUM_TRAINING_IMAGES_EXT))\nNUM_TRAINING_IMAGES_EXT = count_data_items(TRAINING_FILENAMES_IN18, labeled=True)\nprint('IN18; {} images'.format(NUM_TRAINING_IMAGES_EXT))\nNUM_TRAINING_IMAGES_EXT = count_data_items(TRAINING_FILENAMES_IN19, labeled=True)\nprint('IN19; {} images'.format(NUM_TRAINING_IMAGES_EXT))\nNUM_TRAINING_IMAGES_EXT = count_data_items(TRAINING_FILENAMES_OX, labeled=True)\nprint('OX; {} images'.format(NUM_TRAINING_IMAGES_EXT))\nNUM_TRAINING_IMAGES_EXT = count_data_items(TRAINING_FILENAMES_TF, labeled=True)\nprint('TF; {} images'.format(NUM_TRAINING_IMAGES_EXT))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset visualizations"},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAINING_FILENAMES=TRAINING_FILENAMES_IN19","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Peek at training data\ntraining_dataset = get_training_dataset()\ntraining_dataset = training_dataset.unbatch().batch(20)\ntrain_batch = iter(training_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# run this cell again for next set of images\ndisplay_batch_of_images(next(train_batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del training_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# peer at test data\ntest_dataset = get_test_dataset()\ntest_dataset = test_dataset.unbatch().batch(20)\ntest_batch = iter(test_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# run this cell again for next set of images\ndisplay_batch_of_images(next(test_batch))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Keras training\n## Model"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"with strategy.scope():\n    pretrained_model = tf.keras.applications.Xception(weights='imagenet', include_top=False ,input_shape=[*IMAGE_SIZE, 3])\n    pretrained_model.trainable = True # False = transfer learning, True = fine-tuning\n    \n    model = tf.keras.Sequential([\n        pretrained_model,\n        tf.keras.layers.GlobalAveragePooling2D(),\n        tf.keras.layers.Dense(len(CLASSES), activation='softmax')\n    ])\n        \n    model.compile(\n        optimizer='adam',\n        loss = 'sparse_categorical_crossentropy',\n        metrics=['sparse_categorical_accuracy']\n    )\n    \n    lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)\n    \n    model.summary()"},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"start_time = time.time()\n\nhistory = model.fit(get_training_dataset(), steps_per_epoch=STEPS_PER_EPOCH, epochs=EPOCHS,\n                    validation_data=get_validation_dataset(), callbacks=[lr_callback])\n\nkeras_fit_training_time = time.time() - start_time\nprint(\"KERAS FIT TRAINING TIME: {:0.1f}s\".format(keras_fit_training_time))"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"display_training_curves(history.history['loss'], history.history['val_loss'], 'loss', 211)\ndisplay_training_curves(history.history['sparse_categorical_accuracy'], history.history['val_sparse_categorical_accuracy'], 'accuracy', 212)"},{"metadata":{},"cell_type":"markdown","source":"# Custom training loop\n## Model"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"with strategy.scope():\n    pretrained_model = tf.keras.applications.Xception(weights='imagenet', include_top=False ,input_shape=[*IMAGE_SIZE, 3])\n    pretrained_model.trainable = True # False = transfer learning, True = fine-tuning\n    \n    model = tf.keras.Sequential([\n        pretrained_model,\n        tf.keras.layers.GlobalAveragePooling2D(),\n        tf.keras.layers.Dense(len(CLASSES), activation='softmax')\n    ])\n    model.summary()\n    \n    # Instiate optimizer with learning rate schedule\n    class LRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n        def __call__(self, step):\n            return lrfn(epoch=step//STEPS_PER_EPOCH)\n    optimizer = tf.keras.optimizers.Adam(learning_rate=LRSchedule())\n        \n    # this also works but is not very readable\n    # optimizer = tf.keras.optimizers.Adam(learning_rate=lambda: lrfn(tf.cast(optimizer.iterations, tf.float32)//STEPS_PER_EPOCH))\n    \n    # Instantiate metrics\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n    valid_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n    train_loss = tf.keras.metrics.Sum()\n    valid_loss = tf.keras.metrics.Sum()\n    \n    # loss as recommended by the custom training loop Tensorflow documentation:\n    # https://www.tensorflow.org/tutorials/distribute/custom_training#define_the_loss_function\n    # Here, a simpler loss_fn = tf.keras.losses.sparse_categorical_crossentropy would work the same.\n    loss_fn = lambda a,b: tf.nn.compute_average_loss(tf.keras.losses.sparse_categorical_crossentropy(a,b), global_batch_size=BATCH_SIZE)"},{"metadata":{},"cell_type":"markdown","source":"## Step functions"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"@tf.function\ndef train_step(images, labels):\n    with tf.GradientTape() as tape:\n        probabilities = model(images, training=True)\n        loss = loss_fn(labels, probabilities)\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n    \n    # update metrics\n    train_accuracy.update_state(labels, probabilities)\n    train_loss.update_state(loss)\n\n@tf.function\ndef valid_step(images, labels):\n    probabilities = model(images, training=False)\n    loss = loss_fn(labels, probabilities)\n    \n    # update metrics\n    valid_accuracy.update_state(labels, probabilities)\n    valid_loss.update_state(loss)"},{"metadata":{},"cell_type":"markdown","source":"## Training loop"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"start_time = epoch_start_time = time.time()\n\n# distribute the datset according to the strategy\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\nvalid_dist_ds = strategy.experimental_distribute_dataset(get_validation_dataset())\n\nprint(\"Steps per epoch:\", STEPS_PER_EPOCH)\nHistory = namedtuple('History', 'history')\nhistory = History(history={'loss': [], 'val_loss': [], 'sparse_categorical_accuracy': [], 'val_sparse_categorical_accuracy': []})\n\nepoch = 0\nfor step, (images, labels) in enumerate(train_dist_ds):\n    \n    # run training step\n    strategy.experimental_run_v2(train_step, args=(images, labels))\n    print('=', end='', flush=True)\n\n    # validation run at the end of each epoch\n    if ((step+1) // STEPS_PER_EPOCH) > epoch:\n        print('|', end='', flush=True)\n        \n        # validation run\n        for image, labels in valid_dist_ds:\n            strategy.experimental_run_v2(valid_step, args=(image, labels))\n            print('=', end='', flush=True)\n\n        # compute metrics\n        history.history['sparse_categorical_accuracy'].append(train_accuracy.result().numpy())\n        history.history['val_sparse_categorical_accuracy'].append(valid_accuracy.result().numpy())\n        history.history['loss'].append(train_loss.result().numpy() / STEPS_PER_EPOCH)\n        history.history['val_loss'].append(valid_loss.result().numpy() / VALIDATION_STEPS)\n        \n        # report metrics\n        epoch_time = time.time() - epoch_start_time\n        print('\\nEPOCH {:d}/{:d}'.format(epoch+1, EPOCHS))\n        print('time: {:0.1f}s'.format(epoch_time),\n              'loss: {:0.4f}'.format(history.history['loss'][-1]),\n              'accuracy: {:0.4f}'.format(history.history['sparse_categorical_accuracy'][-1]),\n              'val_loss: {:0.4f}'.format(history.history['val_loss'][-1]),\n              'val_acc: {:0.4f}'.format(history.history['val_sparse_categorical_accuracy'][-1]),\n              'lr: {:0.4g}'.format(lrfn(epoch)), flush=True)\n        \n        # set up next epoch\n        epoch = (step+1) // STEPS_PER_EPOCH\n        epoch_start_time = time.time()\n        train_accuracy.reset_states()\n        valid_accuracy.reset_states()\n        valid_loss.reset_states()\n        train_loss.reset_states()\n        \n        if epoch >= EPOCHS:\n            break\n    \nsimple_ctl_training_time = time.time() - start_time\nprint(\"SIMPLE CTL TRAINING TIME: {:0.1f}s\".format(simple_ctl_training_time))"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"display_training_curves(history.history['loss'], history.history['val_loss'], 'loss', 211)\ndisplay_training_curves(history.history['sparse_categorical_accuracy'], history.history['val_sparse_categorical_accuracy'], 'accuracy', 212)"},{"metadata":{},"cell_type":"markdown","source":"# Optimized custom training loop\nOptimized by calling the TPU less often and performing more steps per call\n## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -q efficientnet\nimport efficientnet.tfkeras as efn","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#earlystopping = tf.keras.callbacks.EarlyStopping(monitor='val_sparse_categorical_accuracy', min_delta=0.001, patience=4, verbose=1, mode='auto', baseline=None, restore_best_weights=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    pretrained_model = efn.EfficientNetB0(\n        input_shape=(IMAGE_SIZE[0], IMAGE_SIZE[0], 3),\n        weights='noisy-student',\n        include_top=False\n    )\n\n    #pretrained_model = tf.keras.applications.Xception(weights='imagenet', include_top=False ,input_shape=[*IMAGE_SIZE, 3])\n    pretrained_model.trainable = True # False = transfer learning, True = fine-tuning\n    \n    model = tf.keras.Sequential([\n        pretrained_model,\n        tf.keras.layers.GlobalAveragePooling2D(),\n        tf.keras.layers.Dense(len(CLASSES), activation='softmax')\n    ])\n    model.summary()\n    \n    # Instiate optimizer with learning rate schedule\n    class LRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n        def __call__(self, step):\n            return lrfn(epoch=step//STEPS_PER_EPOCH)\n    optimizer = tf.keras.optimizers.Adam(learning_rate=LRSchedule())\n        \n    # this also works but is not very readable\n    #optimizer = tf.keras.optimizers.Adam(learning_rate=lambda: lrfn(tf.cast(optimizer.iterations, tf.float32)//STEPS_PER_EPOCH))\n    \n    # Instantiate metrics\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n    valid_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n    train_loss = tf.keras.metrics.Sum()\n    valid_loss = tf.keras.metrics.Sum()\n    \n    # Loss\n    # The recommendation from the Tensorflow custom training loop  documentation is:\n    # loss_fn = lambda a,b: tf.nn.compute_average_loss(tf.keras.losses.sparse_categorical_crossentropy(a,b), global_batch_size=BATCH_SIZE)\n    # https://www.tensorflow.org/tutorials/distribute/custom_training#define_the_loss_function\n    # This works too and shifts all the averaging to the training loop which is easier:\n    loss_fn = tf.keras.losses.sparse_categorical_crossentropy","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Step functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"STEPS_PER_TPU_CALL = 99\nVALIDATION_STEPS_PER_TPU_CALL = 29\n\n@tf.function\ndef train_step(data_iter):\n    def train_step_fn(images, labels):\n        with tf.GradientTape() as tape:\n            probabilities = model(images, training=True)\n            loss = loss_fn(labels, probabilities)\n        grads = tape.gradient(loss, model.trainable_variables)\n        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n        \n        #update metrics\n        train_accuracy.update_state(labels, probabilities)\n        train_loss.update_state(loss)\n        \n    # this loop runs on the TPU\n    for _ in tf.range(STEPS_PER_TPU_CALL):\n        strategy.experimental_run_v2(train_step_fn, next(data_iter))\n\n@tf.function\ndef valid_step(data_iter):\n    def valid_step_fn(images, labels):\n        probabilities = model(images, training=False)\n        loss = loss_fn(labels, probabilities)\n        \n        # update metrics\n        valid_accuracy.update_state(labels, probabilities)\n        valid_loss.update_state(loss)\n        \n    # this loop runs on the TPU\n    for _ in tf.range(VALIDATION_STEPS_PER_TPU_CALL):\n        strategy.experimental_run_v2(valid_step_fn, next(data_iter))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"start_time = epoch_start_time = time.time()\n\n# distribute the datset according to the strategy\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\n# Hitting End Of Dataset exceptions is a problem in this setup. Using a repeated validation set instead.\n# This will introduce a slight inaccuracy because the validation dataset now has some repeated elements.\nvalid_dist_ds = strategy.experimental_distribute_dataset(get_validation_dataset(repeated=True))\n\nprint(\"Training steps per epoch:\", STEPS_PER_EPOCH, \"in increments of\", STEPS_PER_TPU_CALL)\nprint(\"Validation images:\", NUM_VALIDATION_IMAGES,\n      \"Batch size:\", BATCH_SIZE,\n      \"Validation steps:\", NUM_VALIDATION_IMAGES//BATCH_SIZE, \"in increments of\", VALIDATION_STEPS_PER_TPU_CALL)\nprint(\"Repeated validation images:\", int_div_round_up(NUM_VALIDATION_IMAGES, BATCH_SIZE*VALIDATION_STEPS_PER_TPU_CALL)*VALIDATION_STEPS_PER_TPU_CALL*BATCH_SIZE-NUM_VALIDATION_IMAGES)\nHistory = namedtuple('History', 'history')\nhistory = History(history={'loss': [], 'val_loss': [], 'sparse_categorical_accuracy': [], 'val_sparse_categorical_accuracy': []})\n\nepoch = 0\ntrain_data_iter = iter(train_dist_ds) # the training data iterator is repeated and it is not reset\n                                      # for each validation run (same as model.fit)\nvalid_data_iter = iter(valid_dist_ds) # the validation data iterator is repeated and it is not reset\n                                      # for each validation run (different from model.fit whre the\n                                      # recommendation is to use a non-repeating validation dataset)\n\nstep = 0\nepoch_steps = 0\nwhile True:\n    \n    # run training step\n    train_step(train_data_iter)\n    epoch_steps += STEPS_PER_TPU_CALL\n    step += STEPS_PER_TPU_CALL\n    print('=', end='', flush=True)\n\n    # validation run at the end of each epoch\n    if (step // STEPS_PER_EPOCH) > epoch:\n        print('|', end='', flush=True)\n        \n        # validation run\n        valid_epoch_steps = 0\n        for _ in range(int_div_round_up(NUM_VALIDATION_IMAGES, BATCH_SIZE*VALIDATION_STEPS_PER_TPU_CALL)):\n            valid_step(valid_data_iter)\n            valid_epoch_steps += VALIDATION_STEPS_PER_TPU_CALL\n            print('=', end='', flush=True)\n\n        # compute metrics\n        history.history['sparse_categorical_accuracy'].append(train_accuracy.result().numpy())\n        history.history['val_sparse_categorical_accuracy'].append(valid_accuracy.result().numpy())\n        history.history['loss'].append(train_loss.result().numpy() / (BATCH_SIZE*epoch_steps))\n        history.history['val_loss'].append(valid_loss.result().numpy() / (BATCH_SIZE*valid_epoch_steps))\n        \n        # report metrics\n        epoch_time = time.time() - epoch_start_time\n        print('\\nEPOCH {:d}/{:d}'.format(epoch+1, EPOCHS))\n        print('time: {:0.1f}s'.format(epoch_time),\n              'loss: {:0.4f}'.format(history.history['loss'][-1]),\n              'accuracy: {:0.4f}'.format(history.history['sparse_categorical_accuracy'][-1]),\n              'val_loss: {:0.4f}'.format(history.history['val_loss'][-1]),\n              'val_acc: {:0.4f}'.format(history.history['val_sparse_categorical_accuracy'][-1]),\n              'lr: {:0.4g}'.format(lrfn(epoch)),\n              'steps/val_steps: {:d}/{:d}'.format(epoch_steps, valid_epoch_steps), flush=True)\n        \n        # set up next epoch\n        epoch = step // STEPS_PER_EPOCH\n        epoch_steps = 0\n        epoch_start_time = time.time()\n        train_accuracy.reset_states()\n        valid_accuracy.reset_states()\n        valid_loss.reset_states()\n        train_loss.reset_states()\n        if epoch >= EPOCHS:\n            break\n\noptimized_ctl_training_time = time.time() - start_time\nprint(\"OPTIMIZED CTL TRAINING TIME: {:0.1f}s\".format(optimized_ctl_training_time))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display_training_curves(history.history['loss'], history.history['val_loss'], 'loss', 211)\ndisplay_training_curves(history.history['sparse_categorical_accuracy'], history.history['val_sparse_categorical_accuracy'], 'accuracy', 212)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Confusion matrix"},{"metadata":{"trusted":true},"cell_type":"code","source":"cmdataset = get_validation_dataset(ordered=True) # since we are splitting the dataset and iterating separately on images and labels, order matters.\nimages_ds = cmdataset.map(lambda image, label: image)\nlabels_ds = cmdataset.map(lambda image, label: label).unbatch()\ncm_correct_labels = next(iter(labels_ds.batch(NUM_VALIDATION_IMAGES))).numpy() # get everything as one batch\ncm_probabilities = model.predict(images_ds)\ncm_predictions = np.argmax(cm_probabilities, axis=-1)\nprint(\"Correct   labels: \", cm_correct_labels.shape, cm_correct_labels)\nprint(\"Predicted labels: \", cm_predictions.shape, cm_predictions)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cmat = confusion_matrix(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)))\nscore = f1_score(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)), average='macro')\nprecision = precision_score(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)), average='macro')\nrecall = recall_score(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)), average='macro')\ncmat = (cmat.T / cmat.sum(axis=1)).T # normalized\ndisplay_confusion_matrix(cmat, score, precision, recall)\nprint('f1 score: {:.3f}, precision: {:.3f}, recall: {:.3f}'.format(score, precision, recall))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"B0: (224x224) ep30: f1 score: 0.929, precision: 0.930, recall: 0.932\n\nB3: (224x224) ep30: f1 score: 0.934, precision: 0.931, recall: 0.940 (maybe restore best weights)\n\nB5: (224x224) ep30: f1 score: 0.937, precision: 0.940, recall: 0.939\n\nB5: (331x331) ep30: f1 score: 0.946, precision: 0.948, recall: 0.949 (21sec/ep)\n\nB7: (512x512) ep30: fail in ep2\n\n# ------------------- test val on ext datasets -------------------------------\n* B0: std train; ep20: \n* in17: \n* in18:\n*  in19:\n* ox:\n* tf:\n* IN17,18,19:\n"},{"metadata":{},"cell_type":"markdown","source":"# Predictions"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_ds = get_test_dataset(ordered=True) # since we are splitting the dataset and iterating separately on images and ids, order matters.\n\nprint('Computing predictions...')\ntest_images_ds = test_ds.map(lambda image, idnum: image)\nprobabilities = model.predict(test_images_ds)\npredictions = np.argmax(probabilities, axis=-1)\nprint(predictions)\n\nprint('Generating submission.csv file...')\ntest_ids_ds = test_ds.map(lambda image, idnum: idnum).unbatch()\ntest_ids = next(iter(test_ids_ds.batch(NUM_TEST_IMAGES))).numpy().astype('U') # all in one batch\nnp.savetxt('submission.csv', np.rec.fromarrays([test_ids, predictions]), fmt=['%s', '%d'], delimiter=',', header='id,label', comments='')\n!head submission.csv","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visual validation"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = get_validation_dataset()\ndataset = dataset.unbatch().batch(20)\nbatch = iter(dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# run this cell again for next set of images\nimages, labels = next(batch)\nprobabilities = model.predict(images)\npredictions = np.argmax(probabilities, axis=-1)\ndisplay_batch_of_images((images, labels), predictions)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Timing summary"},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"KERAS FIT TRAINING TIME: {:0.1f}s\".format(keras_fit_training_time))\nprint(\"SIMPLE CTL TRAINING TIME: {:0.1f}s\".format(simple_ctl_training_time))\nprint(\"OPTIMIZED CTL TRAINING TIME: {:0.1f}s\".format(optimized_ctl_training_time))","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.6"}},"nbformat":4,"nbformat_minor":4}