{"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"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":21154,"databundleVersionId":1243559,"sourceType":"competition"}],"dockerImageVersionId":30299,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"**Created by Sanskar Hasija**\n\n**[JAX+FLAX+TF.DATA] Vision Transformers Tutorial 🚀**\n\n**20 February 2021**\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"#
[JAX+FLAX+TF.DATA] VISION TRANSFORMERS TUTORIAL 🚀
\n##
If you find this notebook useful, support with an upvote👍
","metadata":{}},{"cell_type":"markdown","source":"# Table of Contents\n\n[1. Introduction ](#1)
\n[2. Installations](#2)
\n[3. TPU SETUP](#3)
\n[4. TF.data Pipeline](#4)
\n[5. Image Augumentations](#5)
\n[6. Pre-Processing Data for modelling ](#6)
\n[7. Model Define and Configuration ](#7)
\n[8. Model Training ](#8)
\n[9. Evaluating Predictions ](#9)
\n[10. Test Predictions and Submission ](#10)
","metadata":{}},{"cell_type":"markdown","source":"\n# **
Introduction
**","metadata":{}},{"cell_type":"markdown","source":"## Notebook overview : \n* This notebook is end to end notebook for Petals to the Metal - Flower Classification using state of the art ViT (Vision Transformer) architechure.
\n* **The code is reusable to train Vision Transformer on any Image classification dataset.**\n* The ViT model is an transformer based arcitecture for image classification task.The ViT model was intorduced in the research paper [AN IMAGE IS WORTH 16X16 WORDS:\nTRANSFORMERS FOR IMAGE RECOGNITION AT SCALE](https://arxiv.org/pdf/2010.11929.pdf). \n* Pre-trained on imagenet dataset, the ViT model is fine tuned on the dataset in JAX and Flax.Input dataset piplines for training, validaiton and inference(submission) are built using tf.data api. ","metadata":{}},{"cell_type":"markdown","source":" \n

JAX :

\n

\n\n\n\n* JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research. JAX can automatically differentiate native Python and NumPy code.
\n* JAX uses XLA (Accelerated Linear Algebra) to compile and run your NumPy code on accelerators, like GPUs and TPUs. Compilation happens under the hood by default, with library calls getting just-in-time compiled and executed. JAX even lets just-in-time (JIT) compile Python functions into XLA-optimized kernels using a one-function API.
\n* Compilation and automatic differentiation can be composed arbitrarily, to get maximal performance without having to leave Python.
\n\nCheck out JAX Quickstart guide and documentation - [here](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html).
\n ","metadata":{}},{"cell_type":"markdown","source":" \n

Flax :

\n

\n\n\n\n* Flax is a high-performance neural network library for JAX that is designed for flexibility.
\n* Flax was originally started by engineers and researchers within the Brain Team in Google Research (in close collaboration with the JAX team), and is now developed jointly with the open source community.
\n\nCheck out Flax documentation - [here](https://flax.readthedocs.io/en/latest/overview.html#flax).
\n \n \n","metadata":{}},{"cell_type":"markdown","source":" \n

TF.data

\n

\n\n\n\n* TF.DATA API is used for building efficient input pipelines which can handle large amounts of data and perform complex data transformations.
\n* TF.DATA API has provisions for handling different data formats.
\n\nCheck out tf.data guide and documentation - [here](https://www.tensorflow.org/guide/data).
\n\n","metadata":{}},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"#
Installations
\n","metadata":{}},{"cell_type":"markdown","source":"* ViT official repository by Google-Research is cloned and appended in the system path. Check out the GitHub repo - [here](https://github.com/google-research/vision_transformer).","metadata":{}},{"cell_type":"code","source":"from IPython.display import clear_output\n\n![ -d vision_transformer ] || git clone --depth=1 https://github.com/google-research/vision_transformer\n!cd vision_transformer && git pull\n\n!pip install clu==0.0.2\n!pip install einops\n\n!pip3 install -U jaxlib # Updating to latest versioln of jaxlib\n!pip3 install -U jax # Updating to latest versioln of jax\n!pip3 install -U flax # Updating to latest versioln of flax\nclear_output()","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:42:09.902206Z","iopub.execute_input":"2022-02-19T21:42:09.902552Z","iopub.status.idle":"2022-02-19T21:43:07.338405Z","shell.execute_reply.started":"2022-02-19T21:42:09.902464Z","shell.execute_reply":"2022-02-19T21:43:07.337299Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"#
TPU Setup
\n","metadata":{}},{"cell_type":"markdown","source":"## Setting up TPU for JAX : ","metadata":{}},{"cell_type":"markdown","source":"* This notebook works with all three accelerator TPU, GPU and CPU provided by Kaggle.But using TPU is recommended for faster training and evaluation.","metadata":{}},{"cell_type":"code","source":"import os\nif 'TPU_NAME' in os.environ:\n import requests\n if 'TPU_DRIVER_MODE' not in globals():\n url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'\n resp = requests.post(url)\n TPU_DRIVER_MODE = 1\n from jax.config import config\n config.FLAGS.jax_xla_backend = \"tpu_driver\"\n config.FLAGS.jax_backend_target = os.environ['TPU_NAME']\n print(\"\\033[94m\")\n print(\"TPU DETECTED!\")\n print('Registered TPU:', config.FLAGS.jax_backend_target)\nelse:\n print('No TPU detected.')\nimport jax \nDEVICE_COUNT = len(jax.local_devices())\nTPU = DEVICE_COUNT==8\n\nif TPU:\n print(\"\\033[94m\")\n print(\"8 cores of TPU ( Local devices in Jax ):\")\n print('\\n'.join(map(str,jax.local_devices())))","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:43:07.340262Z","iopub.execute_input":"2022-02-19T21:43:07.340959Z","iopub.status.idle":"2022-02-19T21:43:52.800072Z","shell.execute_reply.started":"2022-02-19T21:43:07.340908Z","shell.execute_reply":"2022-02-19T21:43:52.798913Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Imports : \n","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append('./vision_transformer')\n\n#VISION TRANSFORMER REPO IMPORTS \nfrom vit_jax import utils\nfrom vit_jax import train\nfrom vit_jax import models\nfrom vit_jax import checkpoint\nfrom vit_jax import momentum_clip\nfrom vit_jax import input_pipeline\nfrom vit_jax.configs import common as common_config\nfrom vit_jax.configs import models as models_config\n\n#LIBRARIES IMPORT \nimport os\nimport re \nimport jax\nimport flax\nimport math\nimport tqdm\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport jax.numpy as jnp\nfrom scipy import stats\nimport tensorflow as tf\nimport tensorflow_addons as tfa\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import trange\nfrom sklearn.metrics import f1_score\nfrom sklearn.metrics import recall_score \nfrom sklearn.metrics import precision_score\nfrom sklearn.metrics import confusion_matrix\nwarnings.filterwarnings('ignore')","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:43:52.801431Z","iopub.execute_input":"2022-02-19T21:43:52.801911Z","iopub.status.idle":"2022-02-19T21:44:00.337231Z","shell.execute_reply.started":"2022-02-19T21:43:52.801861Z","shell.execute_reply":"2022-02-19T21:44:00.336064Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"# **
TF.data Pipeline
**\n","metadata":{}},{"cell_type":"markdown","source":"\n
\n","metadata":{}},{"cell_type":"markdown","source":"📌 tf.data API is used for building efficient input pipelines which can handle large amounts of data and perform complex data transformations . tf.data API has provisions for handling different data formats.\n \n📌The tf.data API enables you to build complex input pipelines from simple, reusable pieces. The tf.data API also makes it possible to handle large amounts of data, read from different data formats, and perform complex transformations.\n\n📌The tf.data API introduces a tf.data.Dataset abstraction that represents a sequence of elements, in which each element consists of one or more components. For example, in an image pipeline, an element might be a single training example, with a pair of tensor components representing the image and its label.\n\n**Input Pipeline :**\n\n* tf.data.TFRecordDataset() is used to create an input pipeline for data stored in TFRecord format .\n \n","metadata":{}},{"cell_type":"markdown","source":"**Transformations :**\n\n* The Dataset object can be transformed into a new Dataset by chaining method calls on the tf.data.Dataset object . Some of the transformations which can be applied are Dataset.map() , Dataset.batch() , Dataset.shuffle() , Dataset.prefetch() .\n\n* The Dataset object is a Python iterable which it possible to consume its elements using a for loop .\n\n
","metadata":{}},{"cell_type":"markdown","source":"## Setting up TF.data Pipeline : ","metadata":{}},{"cell_type":"code","source":"IMAGE_SIZE = [512,512]\nif TPU :\n from kaggle_datasets import KaggleDatasets\n GCS_DS_PATH = KaggleDatasets().get_gcs_path('tpu-getting-started')\n GCS_PATH = GCS_DS_PATH + '/tfrecords-jpeg-'+ str(IMAGE_SIZE[0])+\"x\" + str(IMAGE_SIZE[1])\n PATH = GCS_PATH \n \nelse: \n PATH = \"../input/tpu-getting-started/tfrecords-jpeg-\"+ str(IMAGE_SIZE[0])+\"x\" + str(IMAGE_SIZE[1])","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:00.339965Z","iopub.execute_input":"2022-02-19T21:44:00.340362Z","iopub.status.idle":"2022-02-19T21:44:00.775844Z","shell.execute_reply.started":"2022-02-19T21:44:00.340316Z","shell.execute_reply":"2022-02-19T21:44:00.774559Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The data is provided serialized into TFRecords. This is a format convenient for distributing data to each of the TPUs cores. \nThe tf.data input pipeline used in this notebook is a modified version of the wonderful beginner notebook - [Create Your First Submission](https://www.kaggle.com/ryanholbrook/create-your-first-submission)\nof the [Computer Vision](https://www.kaggle.com/learn/computer-vision) course provided by Kaggle.","metadata":{}},{"cell_type":"code","source":"AUTO = tf.data.experimental.AUTOTUNE\nSEED = 12 \nTRAINING_FILENAMES = tf.io.gfile.glob(PATH + '/train/*.tfrec')\nVALIDATION_FILENAMES = tf.io.gfile.glob(PATH + '/val/*.tfrec')\nTEST_FILENAMES = tf.io.gfile.glob(PATH + '/test/*.tfrec') \n\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\n# function for decoding and reshaping image\ndef decode_image(image_data): \n \"\"\"\n Decodes jpeg image and reshapes \n it to (IMAGE_HEIGHT, IMAGE_WIDTH, 3)\n \"\"\"\n image = tf.image.decode_jpeg(image_data, channels=3)\n image = tf.cast(image, tf.float32) \n image = tf.reshape(image, [*IMAGE_SIZE, 3]) \n return image\n\n\n# function for reading labeled tf_record\n# useful for loading train and valid data\ndef read_labeled_tfrecord(example): \n \"\"\"\n Reads labeled tf_record for train/valid\n datasets\n \"\"\"\n LABELED_TFREC_FORMAT = { \n \"image\": tf.io.FixedLenFeature([], tf.string), \n \"class\": tf.io.FixedLenFeature([], tf.int64), \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 \n\n\n# function for reading un_labeled tf_record\n# useful for loading test data \ndef read_unlabeled_tfrecord(example): \n \"\"\"\n Reads unlabeled tf_record for \n test dataset\n \"\"\"\n UNLABELED_TFREC_FORMAT = { \n \"image\": tf.io.FixedLenFeature([], tf.string), \n \"id\": tf.io.FixedLenFeature([], tf.string), \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 \n\n\n#function to load datasets\ndef load_dataset(filenames, labeled=True,ordered = False ): \n \"\"\"\n Loads tf.data.Dataset \n \"\"\"\n ignore_order = tf.data.Options()\n if not ordered:\n ignore_order.experimental_deterministic = False\n dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)\n dataset = dataset.with_options(ignore_order) \n dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, \n num_parallel_calls=AUTO)\n return dataset\n\n# function for counting total items\ndef count_data_items(filenames): \n \"\"\"\n Counts the number of \n data items\n \"\"\"\n n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) \n for filename in filenames]\n return np.sum(n)\n\n\nNUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES)\nNUM_VALIDATION_IMAGES = count_data_items(VALIDATION_FILENAMES)\nNUM_TEST_IMAGES = count_data_items(TEST_FILENAMES)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:00.778109Z","iopub.execute_input":"2022-02-19T21:44:00.778824Z","iopub.status.idle":"2022-02-19T21:44:01.073865Z","shell.execute_reply.started":"2022-02-19T21:44:00.778775Z","shell.execute_reply":"2022-02-19T21:44:01.072707Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Now Datasets are loaded and then mapped to a function that names the respective tensor for further processing.","metadata":{}},{"cell_type":"code","source":"train_ds = load_dataset(TRAINING_FILENAMES, labeled=True)\nvalid_ds = load_dataset(VALIDATION_FILENAMES, labeled=True )\ntest_ds = load_dataset(TEST_FILENAMES, labeled=False, ordered=True)\n\n# function for renaming\ndef idx_name(image,label): \n return dict( image =image ,\n label = label)\n\n# function for renaming\ndef idx_name_test(image, image_id): \n return dict(image = image,\n image_id= image_id)\n\n\ntrain_ds = train_ds.map(idx_name,num_parallel_calls=AUTO)\nvalid_ds = valid_ds.map(idx_name,num_parallel_calls=AUTO)\ntest_ds = test_ds.map(idx_name_test,num_parallel_calls=AUTO)","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:44:01.075436Z","iopub.execute_input":"2022-02-19T21:44:01.0759Z","iopub.status.idle":"2022-02-19T21:44:01.663932Z","shell.execute_reply.started":"2022-02-19T21:44:01.075855Z","shell.execute_reply":"2022-02-19T21:44:01.662917Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"# **
Image Augumentations
**\n","metadata":{}},{"cell_type":"markdown","source":"## Visualizing Differnt Augumentations : ","metadata":{}},{"cell_type":"markdown","source":"\n Data augmentation is a technique through which one can increase the size of the data for the training of the model without adding the new data.
Techniques like padding, cropping, rotating, and flipping are the most common methods that are used over the images to increase the data size.
\n
\nThere are six different augumentations/transformations implemented in this notebook :
\n* Random Brighness
\n* Random Contrast
\n* Random Saturation
\n* Random Crop or Pad
\n* Random Rotate
\n* Sharpness
\n","metadata":{}},{"cell_type":"code","source":"VIS_IMAGE = 4\naug_batch = next(iter(train_ds.batch(VIS_IMAGE).as_numpy_iterator()))\nfor image in aug_batch[\"image\"]: \n seed = (SEED, SEED)\n image = image/ 255\n random_bright = tf.image.stateless_random_brightness(image, max_delta=1.0, \n seed = seed)\n random_contrast = tf.image.stateless_random_contrast(image, 0.2, 2.0,\n seed = seed)\n random_saturation = tf.image.stateless_random_saturation(image, 0.2, 1.0,\n seed = seed)\n random = tf.random.uniform(shape=[],minval=-20, maxval=20).numpy().astype(\"int\")\n random_crop_or_pad = tf.image.resize_with_crop_or_pad(image, \n tf.shape(image).numpy()[0] + random, \n tf.shape(image).numpy()[1] + random)\n random_rotate = tfa.image.rotate(image, tf.constant(tf.random.uniform((1,), \n minval = 0.01,\n maxval = 0.4)))\n sharpness = tfa.image.sharpness(image, 5.1)\n\n plt.figure(figsize = (20, 20))\n \n plt.subplot(1, 7, 1)\n plt.imshow(image)\n plt.axis('off')\n plt.title('Original')\n \n plt.subplot(1, 7, 2)\n plt.imshow(random_bright.numpy())\n plt.title('Random Brightness')\n plt.axis('off')\n \n plt.subplot(1, 7, 3)\n plt.imshow(random_contrast.numpy())\n plt.title('Random Contrast')\n plt.axis('off')\n \n plt.subplot(1, 7, 4)\n plt.imshow(random_saturation.numpy())\n plt.title('Random Saturation')\n plt.axis('off')\n \n plt.subplot(1, 7, 5)\n plt.imshow(random_crop_or_pad.numpy())\n plt.title('Random Crop or Pad')\n plt.axis('off')\n \n plt.subplot(1, 7, 6)\n plt.imshow(random_rotate.numpy())\n plt.title('Random Rotate')\n plt.axis('off')\n \n plt.subplot(1, 7, 7)\n plt.imshow(sharpness.numpy())\n plt.title('Sharpness')","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:01.66552Z","iopub.execute_input":"2022-02-19T21:44:01.665985Z","iopub.status.idle":"2022-02-19T21:44:11.213877Z","shell.execute_reply.started":"2022-02-19T21:44:01.665947Z","shell.execute_reply":"2022-02-19T21:44:11.20849Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#### Difference between normal tf.image functions and stateless ones:\n* There are two sets of random image operations: `tf.image.random*` and `tf.image.stateless_random*`. \n\n* Using `tf.image.random*` operations is strongly discouraged as they use the old RNGs from TF 1.x. (TensorFlow Website)","metadata":{}},{"cell_type":"markdown","source":"## Defining Augumentation function for preprocess : ","metadata":{}},{"cell_type":"markdown","source":"* Transformations are divided into `two groups` with a 50% chance of being applied.\n* The two groups are further divided into `two sub-groups` with an equal chance of being applied.\n* Each of the `four sub-groups` transformations have a 25% chance of being applied ","metadata":{}},{"cell_type":"code","source":"def aug(data):\n \"\"\"\n Augument the image with one of four \n different equal likely transformations.\n Transformations:\n 1. Random Contrast\n 2. Random Brightness\n 3. Random Saturation\n 4. Random Rotation\n \"\"\"\n image = data[\"image\"]\n seed = (SEED,SEED)\n \n # two groups for transformations\n transformation_selection = tf.random.uniform([], minval = 0, \n maxval = 1, \n dtype = tf.float32)\n \n # probability for sub groups\n # each transformation has 25% chance\n # of being applied on image\n prob_1 = tf.random.uniform([], minval = 0, maxval = 1, dtype = tf.float32) \n prob_2 = tf.random.uniform([], minval = 0, maxval = 1, dtype = tf.float32)\n image = tf.cond(tf.greater(transformation_selection, 0.5),\n lambda: tf.cond(tf.greater(prob_1, 0.5),\n lambda: tf.image.stateless_random_contrast(image, 0.1, 0.5,\n seed = seed), \n lambda: tf.image.stateless_random_brightness(image, max_delta=0.3, \n seed = seed),\n ),\n lambda: tf.cond(tf.greater(prob_2, 0.5),\n lambda: tf.image.stateless_random_saturation(image, 0.01, 0.1,\n seed = seed),\n lambda: tfa.image.rotate(image, tf.random.uniform((1,), \n minval = 0.01,\n maxval = 0.2)) \n ) \n )\n data[\"image\"] = image\n return data","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:11.221711Z","iopub.execute_input":"2022-02-19T21:44:11.222795Z","iopub.status.idle":"2022-02-19T21:44:11.275157Z","shell.execute_reply.started":"2022-02-19T21:44:11.222669Z","shell.execute_reply":"2022-02-19T21:44:11.273964Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"# **
Pre-Processing Data for modelling
**\n","metadata":{}},{"cell_type":"markdown","source":"## Loading Train and Validation dataset : ","metadata":{}},{"cell_type":"markdown","source":"* The `input_pipline` moduled imported earlier helps us in setting the innput_data pipline.\n* The `train_ds` and `valid_ds` datasets are instantiated using the `get_data` function from `input_pipline` module
\n
\n* The arguments of get_data function are as follows: \n * data: tf.dataset to read data from\n * mode: train or test. `test for evaluation(validation) task`\n * num_classes : Number of output classes (All classes get one-hot encoded) \n * image_decoder : image_decoder function if image is not decoded before passing to the function\n * repeats : Number of time dataset to be repeated\n * batch_size : Batch Size\n * image_size : Image size after cropping (for training) / resizing (for\n evaluation).\n * shuffle_buffer : Buffer size for shuffling dataset\n * preproces : Optional preprocessing for data\n ","metadata":{}},{"cell_type":"markdown","source":"The dataset passed on to the `get_data` function gets augumented according to the preproces function passed, and then processed as required by the model and get reshaped as:
\n####
(IMAGE HEIGHT, IMAGE WIDTH , 3) -->(DEVICE COUNT, BATCH, IMAGE HEIGHT, IMAGE WIDTH , 3)
\n\n* In case of `TPU - DEVICE COUNT = 8`, and in `GPU and CPU DEVICE COUNT = 1`.","metadata":{}},{"cell_type":"code","source":"AUTOTUNE = tf.data.experimental.AUTOTUNE\nNUM_CLASSES = 102\nSHUFFLE = 1024\nIMAGE_DECODER = lambda image : image\nBATCH_SIZE = 16 * DEVICE_COUNT\nIMAGE_SIZE = 512\nPREPROCESS = aug\n\ntrain_ds = input_pipeline.get_data(\n data = train_ds,\n mode = \"train\",\n num_classes = NUM_CLASSES,\n image_decoder = IMAGE_DECODER,\n repeats = None,\n batch_size = BATCH_SIZE,\n image_size = IMAGE_SIZE,\n shuffle_buffer = SHUFFLE,\n preprocess = PREPROCESS\n )\n\n# No augumentation required for validation_dataset\nvalid_ds = input_pipeline.get_data(\n data = valid_ds,\n mode = \"test\",\n num_classes = NUM_CLASSES,\n image_decoder = IMAGE_DECODER,\n repeats = 1,\n batch_size = BATCH_SIZE,\n image_size = IMAGE_SIZE,\n shuffle_buffer = SHUFFLE,\n preprocess = None \n )","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:44:11.27709Z","iopub.execute_input":"2022-02-19T21:44:11.27776Z","iopub.status.idle":"2022-02-19T21:44:13.896402Z","shell.execute_reply.started":"2022-02-19T21:44:11.277715Z","shell.execute_reply":"2022-02-19T21:44:13.894573Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Loading Test dataset : ","metadata":{}},{"cell_type":"markdown","source":"* Unfortunately get_data function from input_pipline does not generate data from inference data without any labels.\n* `test_ds` have been manually processed and `shared across the 8 devices of TPU` (if present)","metadata":{}},{"cell_type":"markdown","source":" \n* Images in test_ds are first normalized, then batched together in batch_size defined before.
\n* After batching the shape changes as :
\n

(IMAGE HEIGHT, IMAGE WIDTH , 3) -->(BATCH, IMAGE HEIGHT, IMAGE WIDTH , 3)

\n * Then batched_images are shared across 8 TPU cores (if present). The shape changes as :
\n

(BATCH, IMAGE HEIGHT, IMAGE WIDTH , 3) -->(DEVICE COUNT, BATCH, IMAGE HEIGHT, IMAGE WIDTH , 3)


\n","metadata":{}},{"cell_type":"code","source":"# Extracting test ids for submission later\ntest_ids_ds = test_ds.map(lambda data : data[\"image_id\"])\ntest_ids = next(iter(test_ids_ds.batch(NUM_TEST_IMAGES))).numpy().astype('U') \n\n\n#Reshaping data for devices \n#(8 in case of TPU and 1\n# in case of GPU and CPU)\ndef shared_data(data):\n \"\"\"\n Reshapes image from (BATCH, IMAGE HEIGHT, IMAGE WIDTH , 3) to \n (DEVICE COUNT, BATCH, IMAGE HEIGHT, IMAGE WIDTH , 3) \n \"\"\"\n data = tf.reshape(data,shape = [DEVICE_COUNT, -1, IMAGE_SIZE, IMAGE_SIZE, 3]) \n return data \n\n \ndef normalize(image):\n \"\"\"\n Normalized an image tensor\n \"\"\"\n image = (image - 127.5) / 127.5\n return image\n\ndef test_data_prep(dataset): \n \"\"\"\n Prepares test data for inference \n similar to get_data \n function of input_pipeline\n \"\"\"\n dataset = dataset.map(lambda data: data[\"image\"])\n dataset = dataset.map(normalize)\n dataset = dataset.batch(BATCH_SIZE, drop_remainder=True) \n dataset = dataset.map(shared_data, tf.data.experimental.AUTOTUNE)\n return dataset.prefetch(1)\n\ntest_ds = test_data_prep(test_ds)","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:44:13.930482Z","iopub.execute_input":"2022-02-19T21:44:13.933528Z","iopub.status.idle":"2022-02-19T21:44:30.751791Z","shell.execute_reply.started":"2022-02-19T21:44:13.933465Z","shell.execute_reply":"2022-02-19T21:44:30.750552Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualizing Processed Training Images : ","metadata":{"execution":{"iopub.status.busy":"2022-02-18T15:35:12.642433Z","iopub.execute_input":"2022-02-18T15:35:12.642815Z","iopub.status.idle":"2022-02-18T15:35:12.650104Z","shell.execute_reply.started":"2022-02-18T15:35:12.642779Z","shell.execute_reply":"2022-02-18T15:35:12.649365Z"}}},{"cell_type":"code","source":"def show_image_grid(images, labels ):\n \"\"\"\n Create a plot of 32 images\n \"\"\"\n rows = 4\n cols = 8\n fig = plt.figure(figsize=(20,11))\n for i in range(1, cols*rows+1 ):\n img = images[i-1]\n img = (img +1) / 2 \n fig.add_subplot(rows, cols, i)\n plt.axis('off')\n plt.title(CLASSES[labels[i-1]], fontsize=14)\n plt.imshow(img)\n plt.subplots_adjust(wspace=0.05, hspace=0.05)\n plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:30.753385Z","iopub.execute_input":"2022-02-19T21:44:30.753798Z","iopub.status.idle":"2022-02-19T21:44:30.763877Z","shell.execute_reply.started":"2022-02-19T21:44:30.753758Z","shell.execute_reply":"2022-02-19T21:44:30.762562Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# You need to change this code if you change batch_size\ntrain_batch = next(iter(train_ds))\nimages = np.append(\n train_batch['image'][0].numpy(), \n train_batch['image'][1].numpy(), \n axis = 0)\nlabels = np.append(\n train_batch['label'][0].numpy().argmax(axis = 1), \n train_batch['label'][0].numpy().argmax(axis = 1))\nshow_image_grid(images, labels)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:30.765393Z","iopub.execute_input":"2022-02-19T21:44:30.766022Z","iopub.status.idle":"2022-02-19T21:44:42.326477Z","shell.execute_reply.started":"2022-02-19T21:44:30.765979Z","shell.execute_reply":"2022-02-19T21:44:42.325139Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualizing Validation Images : ","metadata":{}},{"cell_type":"code","source":"# You need to change this code if you change batch_size\nvalid_batch = next(iter(valid_ds))\nimages = np.append(\n valid_batch['image'][0].numpy(), \n valid_batch['image'][1].numpy(), \n axis = 0)\nlabels = np.append(\n valid_batch['label'][0].numpy().argmax(axis = 1) , \n valid_batch['label'][0].numpy().argmax(axis = 1))\nshow_image_grid(images, labels)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:42.328105Z","iopub.execute_input":"2022-02-19T21:44:42.328565Z","iopub.status.idle":"2022-02-19T21:44:48.613958Z","shell.execute_reply.started":"2022-02-19T21:44:42.328517Z","shell.execute_reply":"2022-02-19T21:44:48.612815Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"\n# **
Model Define and Configuration
**","metadata":{}},{"cell_type":"markdown","source":"## Vision Transformer Architecture : ","metadata":{}},{"cell_type":"markdown","source":"
","metadata":{}},{"cell_type":"markdown","source":"* Transformers solve a problem that was not only limited to NLP, but also to Computer Vision tasks.\n* Huge models (ViT-H) generally do better than large models (ViT-L) and win against state-of-the-art methods.\n* Vision transformers work better on large-scale data.\n* Attention Rollouts are used to compute the attention maps.\n* Like the GPT-3 and BERT models, the Visual Transformer model also can scale.\n* Large-scale training outperforms inductive bias.\n* Convolutions are translation invariant, locality-sensitive, and lack a global understanding of images","metadata":{}},{"cell_type":"markdown","source":"Link to the paper - AN IMAGE IS WORTH 16X16 WORDS TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE - [here](https://arxiv.org/pdf/2010.11929.pdf). ","metadata":{}},{"cell_type":"markdown","source":"## Downloading model : ","metadata":{}},{"cell_type":"markdown","source":"#### All available pre-trained models: ","metadata":{}},{"cell_type":"code","source":"!gsutil ls -lh gs://vit_models/imagenet*","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:48.615634Z","iopub.execute_input":"2022-02-19T21:44:48.615973Z","iopub.status.idle":"2022-02-19T21:44:54.501Z","shell.execute_reply.started":"2022-02-19T21:44:48.615939Z","shell.execute_reply":"2022-02-19T21:44:54.499333Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"All the models can also be downloaded from Google cloud Storage platfrom from [here](https://console.cloud.google.com/storage/browser/vit_models/imagenet21k%2Bimagenet2012?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22).\n
\n","metadata":{}},{"cell_type":"markdown","source":"#### Downloading VIT_B_32 model : ","metadata":{}},{"cell_type":"markdown","source":"* `VIT_B_32` model pretrained on `imagenet21k` and `imagenet2012` is used in this notebook for training.","metadata":{}},{"cell_type":"code","source":"model_name = 'ViT-B_32' \n![ -e \"$model_name\".npz ] || gsutil cp gs://vit_models/imagenet21k+imagenet2012/\"$model_name\".npz .\nassert os.path.exists(f'{model_name}.npz')","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:44:54.50336Z","iopub.execute_input":"2022-02-19T21:44:54.504238Z","iopub.status.idle":"2022-02-19T21:45:03.8959Z","shell.execute_reply.started":"2022-02-19T21:44:54.504175Z","shell.execute_reply":"2022-02-19T21:45:03.894705Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Configuration : ","metadata":{}},{"cell_type":"markdown","source":"* `model_config` can be fetched used `MODEL_CONFIG`function from the `models_config` module.\n* `model_config` contains the configuration of the given model.\n* After loading model definition(conifg), `random parameters` are initialized using `VisionTransformer` function from the `models` module.","metadata":{}},{"cell_type":"code","source":"model_name = 'ViT-B_32' \nmodel_config = models_config.MODEL_CONFIGS[model_name]\n\n# Intilizaing model\nmodel = models.VisionTransformer(num_classes=NUM_CLASSES, **model_config) \nprint(f'\\033[94mMODEL Configuration : ')\nprint(model_config)","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:45:03.897696Z","iopub.execute_input":"2022-02-19T21:45:03.898062Z","iopub.status.idle":"2022-02-19T21:45:03.908471Z","shell.execute_reply.started":"2022-02-19T21:45:03.898019Z","shell.execute_reply":"2022-02-19T21:45:03.907501Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Loading Pre-trained model : ","metadata":{}},{"cell_type":"markdown","source":"* Using jax.jit() model is complied to Accelerated Linear Algebra (XLA compiler)\n\n### `jax.jit()`:\n* JAX provides a JIT (just-in-time) compiler which takes a standard Python/ NumPy function and compiles it to run efficiently on an accelerator. Compiling a function also avoids the overhead of the Python interpreter, which helps whether or not you're using an accelerator.\n* XLA(Accelerated Linear Algebra) is used to just-in-time (JIT)-compile and execute JAX programs on GPU and Cloud TPU accelerators.\n\n**Learn more about `jax_jit()` [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit).**","metadata":{}},{"cell_type":"markdown","source":"### `params`:\n* The checkpoint of the pretrained model ( which was downloaded earlier) is loaded using the `load_pretrained` function from the 'checkpoint' module.\n* The pretrained model is loaded and slight changes are made to the parameters of the model like :\n * Changing the final layers.\n * Resizing the position embedding.","metadata":{}},{"cell_type":"code","source":"#Initiliazing random parameters \nvariables= jax.jit(lambda: model.init(jax.random.PRNGKey(SEED), \n train_batch['image'][0, :1],\n train=False),\n backend='cpu')()\n\n#Loading pretrained checkpoint\nparams = checkpoint.load_pretrained(pretrained_path=f'{model_name}.npz', \n init_params=variables['params'],\n model_config=model_config,\n)","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:45:03.909968Z","iopub.execute_input":"2022-02-19T21:45:03.910599Z","iopub.status.idle":"2022-02-19T21:47:19.190098Z","shell.execute_reply.started":"2022-02-19T21:45:03.910554Z","shell.execute_reply":"2022-02-19T21:47:19.188916Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Replicating parameters : ","metadata":{}},{"cell_type":"markdown","source":"* Using `flax.jax_utils.replicate()` every array is replicated to multiple devices.\n* Every array in the pytree is `replicated` into the devices which makes `same data` replicated across `all local devices`.\n* `flax.jax_utils.replicate()` on different devices:\n * TPU : replicates the params in every core. (TOTAL 8 CORES )\n * GPU : moves the data onto the GPU device.\n * CPU : creates a copy.\n\n**Learn more about `lax.jax_utils.replicate()` [here](https://flax.readthedocs.io/en/latest/flax.jax_utils.html).**","metadata":{}},{"cell_type":"markdown","source":"* After replication array type and shape changes as:\n

DeviceArray ->_ShardedDeviceArray

\n

(NUM_CLASSES) ->(DEVICE_COUNT,NUM_CLASSES)


\n(NOTE: The above shape transformation demonstration is of the last output array of params)\n","metadata":{}},{"cell_type":"code","source":"params_repl = flax.jax_utils.replicate(params) ","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:47:19.192257Z","iopub.execute_input":"2022-02-19T21:47:19.192735Z","iopub.status.idle":"2022-02-19T21:47:19.558434Z","shell.execute_reply.started":"2022-02-19T21:47:19.192691Z","shell.execute_reply":"2022-02-19T21:47:19.557183Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\033[94m\")\nprint(\"Params(before replicaiton)- Type: \" \n + str( type(params['head']['bias']).__name__) + \n str(\" \") +\"Shape : \" + \n str(params['head']['bias'].shape))\nprint(\"Params(after replication)- Type: \" + \n str(type(params_repl['head']['bias']).__name__) +\n \" Shape : \" + \n str(params_repl['head']['bias'].shape))","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:47:19.56031Z","iopub.execute_input":"2022-02-19T21:47:19.560992Z","iopub.status.idle":"2022-02-19T21:47:19.572261Z","shell.execute_reply.started":"2022-02-19T21:47:19.560947Z","shell.execute_reply":"2022-02-19T21:47:19.570272Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Mapping model's forward pass : ","metadata":{}},{"cell_type":"markdown","source":"### `jax.pmap`:\n* Applying `pmap()` to a function will compile the function with XLA (similarly to jit()), then execute it in parallel on XLA devices, such as multiple GPUs or multiple TPU cores. \n* `pmap` is different than `vmap` in how values are computer. `vmap` vectorizes a function by adding a batch dimension to every primitive operation in the function, whereas `pmap` replicates the function and executes each replica on its own XLA device in parallel.\n\n**Learn more about `jax.pmap()` [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html#jax.pmap).**
\n**Learn more about `jax.vmap()` [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html#jax.vmap).**","metadata":{}},{"cell_type":"code","source":"vit_apply_repl = jax.pmap(\n lambda params, inputs: model.apply(dict(params=params), \n inputs, \n train=False))","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:47:19.574693Z","iopub.execute_input":"2022-02-19T21:47:19.575223Z","iopub.status.idle":"2022-02-19T21:47:19.585351Z","shell.execute_reply.started":"2022-02-19T21:47:19.575156Z","shell.execute_reply":"2022-02-19T21:47:19.584195Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#
Model Training
\n","metadata":{}},{"cell_type":"markdown","source":"## Utility Functions: ","metadata":{}},{"cell_type":"code","source":"def learning_rate_scheduler(base_learning_rate,total_steps,warmup_steps):\n \"\"\"\n Cosine decay learning_rate_scheduler\n \"\"\"\n def step_fn(step):\n progress = (step - warmup_steps) / float(total_steps - warmup_steps) \n #clip the value between 0 and 1 \n progress = jnp.clip(progress, 0.0, 1.0) \n lr = base_learning_rate * 0.5 * (1. + jnp.cos(jnp.pi * progress)) \n if warmup_steps:\n lr = lr * jnp.minimum(1., step / warmup_steps)\n return jnp.asarray(lr, dtype=jnp.float32)\n return step_fn\n\ndef save_best_params(current_opt,best_opt ,current_loss, best_loss):\n \"\"\"\n Return best params on the basis of decrement in value of loss\n \"\"\"\n if current_loss < best_loss:\n return (current_opt,current_loss)\n else:\n return (best_opt,best_loss)\n \ndef evaluate(params, dataset, image_count , batch_size):\n \"\"\"\n Evaluates the model on the \n validation dataset and returns \n accuracy and preditions and\n true labels\n \"\"\"\n total = 0 \n correct = 0 \n steps = image_count // batch_size\n preds = jnp.array([])\n true_labels = jnp.array([])\n for _, batch in zip(trange(steps), dataset.as_numpy_iterator()):\n prediction = vit_apply_repl(params, batch['image']) \n batch_prediction = prediction.argmax(axis=-1).reshape(BATCH_SIZE )\n batch_true = batch['label'].argmax(axis=-1).reshape(BATCH_SIZE )\n preds = jnp.append(preds,batch_prediction)\n true_labels = jnp.append(true_labels,batch_true)\n preds_bool = prediction.argmax(axis=-1) == batch['label'].argmax(axis=-1) \n correct += preds_bool.sum()\n total += len(preds_bool.flatten())\n accuracy = correct / total\n return accuracy, preds , true_labels\n\ndef predictions(params, dataset, image_count , batch_size):\n \"\"\"\n Return predictions by model \n on test dataset for inference\n \"\"\"\n steps = image_count // batch_size \n preds = jnp.array([])\n for s, batch in zip(trange(steps), dataset.as_numpy_iterator()):\n prediction = vit_apply_repl(params, batch) \n batch_prediction = prediction.argmax(axis=-1).reshape(BATCH_SIZE )\n preds = jnp.append(preds,batch_prediction)\n return preds\n\n## FUNCTION TAKEN FROM @ryanholbrook tutorial notebook.\n## https://www.kaggle.com/ryanholbrook/create-your-first-submission\ndef display_confusion_matrix(cmat, score, precision, recall):\n \"\"\"\n Plots Confusion Matrix \n \"\"\"\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, \n fontdict={\n 'fontsize': 18, \n 'horizontalalignment':'right',\n 'verticalalignment':'top', \n 'color':'#804040'})\n plt.show()\n ","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:47:19.587754Z","iopub.execute_input":"2022-02-19T21:47:19.588064Z","iopub.status.idle":"2022-02-19T21:47:19.62243Z","shell.execute_reply.started":"2022-02-19T21:47:19.588027Z","shell.execute_reply":"2022-02-19T21:47:19.620805Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Cosine Decay Learning Rate Scheduler: ","metadata":{}},{"cell_type":"code","source":"TOTAL_STEPS = 500 # total training steps \nBASE_LR= 0.01 # Base learning rate \nWARMUP_STEPS = 100 # Warmup steps for learning_rate_scheduler\nGRAD_NORM_CLIP = 1 # Glips gradient norm \nACCUM_STEPS = 8 # Controls in how many forward passes the batch is split (8 for TPU)\nVERBOSE = 50 # Verbosity\n\nlr_function = learning_rate_scheduler(base_learning_rate = BASE_LR, \n total_steps = TOTAL_STEPS,\n warmup_steps = WARMUP_STEPS)\n #check utlilty function for details \nplot_lr = []\nfor step in range(TOTAL_STEPS):\n plot_lr.append(lr_function(step).item())\nplt.plot(plot_lr)\nplt.xlabel(\"Steps\")\nplt.ylabel(\"Learning Rate\")\nplt.title(\"Learning Rate Cosine Decay\");","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:47:19.624359Z","iopub.execute_input":"2022-02-19T21:47:19.624826Z","iopub.status.idle":"2022-02-19T21:47:29.010316Z","shell.execute_reply.started":"2022-02-19T21:47:19.624776Z","shell.execute_reply":"2022-02-19T21:47:29.009223Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training : ","metadata":{}},{"cell_type":"code","source":" # Update step for data parallel training\nupdate_fn_repl = train.make_update_fn(apply_fn=model.apply, \n accum_steps= ACCUM_STEPS , \n lr_fn=lr_function)\n\n# Momentum optimizer that stores state using half-precision\nopt = momentum_clip.Optimizer(grad_norm_clip=GRAD_NORM_CLIP).create(params) \n\n# Replicating optimizer across devices\nopt_repl = flax.jax_utils.replicate(opt) \n\n# Initialize PRNGs for dropout.\nupdate_rng_repl = flax.jax_utils.replicate(jax.random.PRNGKey(SEED)) ","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T21:47:29.011975Z","iopub.execute_input":"2022-02-19T21:47:29.01278Z","iopub.status.idle":"2022-02-19T21:47:31.291028Z","shell.execute_reply.started":"2022-02-19T21:47:29.012736Z","shell.execute_reply":"2022-02-19T21:47:31.289915Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"* I have implemented a custom train function for ease and simplicity.\n* The function provides an option to get the best parameter based on the minimum training loss.\n* Verbosity can also be set manually.\n* Option to plot `loss vs steps graph` at the end of training.\n* See the docstring for more details","metadata":{}},{"cell_type":"code","source":"def fit(opt_repl,\n update_rng_repl, \n total_steps, \n data, \n save_best_parameters=True, \n verbose = 0 , \n plot_loss_graph = True ):\n \"\"\"\n Returns paramters after training\n Args:\n opt_repl : Optimier replicated across devices (8 devices for TPU).\n update_rng_rel : PRGs for dropout.\n total_steps : Int, Total training steps.\n data : Data for training.\n save_best_parameters : Bool, saves best parameters based on decreasing loss.\n verbose : Int, Verbosity of training \n plot_loss_graph : Bool, plots loss vs steps graph at the end of training.\n \"\"\"\n losses = []\n if save_best_parameters:\n best_loss = float('inf')\n best_opt = 0 \n \n for step, batch in zip(trange(1, total_steps + 1), data.as_numpy_iterator()):\n opt_repl, loss_repl, update_rng_repl = update_fn_repl(\n opt_repl, \n flax.jax_utils.replicate(step), \n batch, \n update_rng_repl\n )\n \n if save_best_parameters:\n best_opt , best_loss = save_best_params(\n current_opt = opt_repl,\n best_opt = best_opt ,\n current_loss = loss_repl[0],\n best_loss = best_loss\n )\n \n if verbose:\n if step%verbose == 0 or step == 1:\n if save_best_parameters:\n print(f\"\\033[94mSTEP: {step}, LOSS: {loss_repl[0]:.2f}, MIN LOSS: {best_loss:.2f}\")\n else : \n print(f\"\\033[94mSTEP: {step}, LOSS: {loss_repl[0]:.2f}\")\n \n losses.append(loss_repl[0])\n \n if plot_loss_graph:\n plt.plot(losses)\n plt.xlabel(\"Steps\")\n plt.ylabel(\"Loss\")\n plt.title(\"Training Loss\")\n \n if save_best_parameters:\n return best_opt\n return opt_repl","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:47:31.292841Z","iopub.execute_input":"2022-02-19T21:47:31.293199Z","iopub.status.idle":"2022-02-19T21:47:31.308542Z","shell.execute_reply.started":"2022-02-19T21:47:31.293157Z","shell.execute_reply":"2022-02-19T21:47:31.307185Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"opt_repl = fit(opt_repl = opt_repl ,\n update_rng_repl = update_rng_repl,\n total_steps = TOTAL_STEPS,\n data = train_ds,\n save_best_parameters=True, \n verbose = VERBOSE, \n plot_loss_graph = True\n )","metadata":{"execution":{"iopub.status.busy":"2022-02-19T21:47:31.3105Z","iopub.execute_input":"2022-02-19T21:47:31.310828Z","iopub.status.idle":"2022-02-19T22:03:34.264901Z","shell.execute_reply.started":"2022-02-19T21:47:31.310786Z","shell.execute_reply":"2022-02-19T22:03:34.263585Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"#
Evaluating Predictions
\n","metadata":{}},{"cell_type":"code","source":"valid_acc_jnp, valid_preds_jnp, valid_true_jnp = evaluate(\n opt_repl.target, \n valid_ds , \n NUM_VALIDATION_IMAGES , \n BATCH_SIZE\n)\n\nval_acc = valid_acc_jnp.item()*100\nprint(f\"Validation Accuracy = {val_acc:.2f}%\")","metadata":{"execution":{"iopub.status.busy":"2022-02-19T22:03:34.26678Z","iopub.execute_input":"2022-02-19T22:03:34.267659Z","iopub.status.idle":"2022-02-19T22:04:24.25801Z","shell.execute_reply.started":"2022-02-19T22:03:34.267593Z","shell.execute_reply":"2022-02-19T22:04:24.256934Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Confusion Matrix with various evaluation metrics : ","metadata":{}},{"cell_type":"code","source":"y_true = valid_true_jnp.astype(\"int32\")\ny_preds = valid_preds_jnp.astype(\"int32\")\n\nlabels = range(len(CLASSES))\ncmat = confusion_matrix(y_true, y_preds, labels= labels)\nscore = f1_score(y_true, y_preds, labels=labels, average='macro')\nprecision = precision_score(y_true, y_preds, labels=labels, average='macro')\nrecall = recall_score(y_true, y_preds, labels=labels, average='macro')\ndisplay_confusion_matrix(cmat,score,precision,recall)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T22:04:24.259893Z","iopub.execute_input":"2022-02-19T22:04:24.260634Z","iopub.status.idle":"2022-02-19T22:04:40.357155Z","shell.execute_reply.started":"2022-02-19T22:04:24.260591Z","shell.execute_reply":"2022-02-19T22:04:40.355817Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"⬆️Back to Table of Contents ⬆️","metadata":{}},{"cell_type":"markdown","source":"#
Test Predictions and Submission
\n","metadata":{}},{"cell_type":"markdown","source":"## Making Predictions : ","metadata":{}},{"cell_type":"code","source":"test_preds = predictions(opt_repl.target, test_ds, NUM_TEST_IMAGES, BATCH_SIZE)","metadata":{"execution":{"iopub.status.busy":"2022-02-19T22:04:40.362536Z","iopub.execute_input":"2022-02-19T22:04:40.36332Z","iopub.status.idle":"2022-02-19T22:06:06.857796Z","shell.execute_reply.started":"2022-02-19T22:04:40.363266Z","shell.execute_reply":"2022-02-19T22:06:06.856603Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Predictions to Submission: ","metadata":{}},{"cell_type":"code","source":"test_preds = np.append(test_preds, \n np.ones(NUM_TEST_IMAGES%BATCH_SIZE)*stats.mode(test_preds)[0].item()).astype(\"int\") \n#unbatched images are predicted as most common prediction","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-02-19T22:06:06.860429Z","iopub.execute_input":"2022-02-19T22:06:06.861114Z","iopub.status.idle":"2022-02-19T22:06:06.870314Z","shell.execute_reply.started":"2022-02-19T22:06:06.861063Z","shell.execute_reply":"2022-02-19T22:06:06.869272Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission = pd.DataFrame(\n data = np.array([test_ids, test_preds ]).T, \n columns = [\"id\", \"label\"]\n)\nsubmission.to_csv(\"submission.csv\", index = False)\nsubmission.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-19T22:06:06.871844Z","iopub.execute_input":"2022-02-19T22:06:06.872149Z","iopub.status.idle":"2022-02-19T22:06:06.967197Z","shell.execute_reply.started":"2022-02-19T22:06:06.872111Z","shell.execute_reply":"2022-02-19T22:06:06.96611Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## References : ","metadata":{}},{"cell_type":"markdown","source":"https://arxiv.org/pdf/2010.11929.pdf\n\nhttps://colab.research.google.com/github/google-research/vision_transformer/blob/master/vit_jax.ipynb\n\nhttps://github.com/google-research/vision_transformer\n\nhttps://www.kaggle.com/ryanholbrook/create-your-first-submission\n\nhttps://www.kaggle.com/usharengaraju/devfest2020identifyingmelanomainlesionimages/\n\nhttps://www.kaggle.com/frightera/tf-data-ensemble-transfer-learning-f1-0-95\n\nhttps://www.kaggle.com/jalammar/intro-to-data-input-pipelines-with-tf-data","metadata":{}},{"cell_type":"markdown","source":" \n \n###
Thank you for reading🙂

\n###
If you have any feedback, please let me know!

\n","metadata":{}}]}