{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Intro to JAX\n[JAX](https://github.com/google/jax) is a framework which is used for high-performance numerical computing and machine learning research developed at [Google Research](https://research.google/) teams. It allows you to build Python applications with a NumPy-consistent API that specializes in differentiating, vectorizing, parallelizing, and compiling to GPU/TPU Just-In-Time. JAX was designed with performance and speed as a first priority, and is natively compatible with common machine learning accelerators such as [GPUs](https://www.kaggle.com/docs/efficient-gpu-usage) and [TPUs](https://www.kaggle.com/docs/tpu). Large ML models can take ages to train -- you might be interested in using JAX for applications where speed and performance are particularly important!\n### When to use JAX vs TensorFlow?\n[TensorFlow](https://www.tensorflow.org/guide) is a fantastic product, with a rich and fully-featured ecosystem, capable of supporting most every use case a machine learning practitioner might have (e.g. [TFLite](https://www.tensorflow.org/lite) for on-device inference computing, [TFHub](https://tfhub.dev/) for sharing pre-trained models, and many additional specialized applications as well). This type of broad mandate both contrasts and compliments JAX's philosophy, which is more narrowly focused on speed and performance.  We recommend using JAX in situations where you do want to maximize speed and performance but you do not require any of the long tail of features and additional functionalities that only the [TensorFlow ecosystem](https://www.tensorflow.org/learn) can provide.\n### Intro to the FLAX\nJust like [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) focuses on speed, other members of the JAX ecosystem are encouraged to specialize as well.  For example, [Flax](https://flax.readthedocs.io/en/latest/) focuses on neural networks and [jgraph](https://github.com/deepmind/jraph) focuses on graph networks.  \n\n[Flax](https://flax.readthedocs.io/en/latest/) is a JAX-based neural network library that was initially developed by  Google Research's Brain Team (in close collaboration with the JAX team) but is now open source.  If you want to train machine learning models on GPUs and TPUs at an accelerated speed, or if you have an ML project that might benefit from bringing together both [Autograd](https://github.com/hips/autograd) and [XLA](https://www.tensorflow.org/xla), consider using [Flax](https://flax.readthedocs.io/en/latest/) for your next project! [Flax](https://flax.readthedocs.io/en/latest/) is especially well-suited for projects that use large language models, and is a popular choice for cutting-edge [machine learning research](https://arxiv.org/search/?query=JAX&searchtype=all&abstracts=show&order=-announced_date_first&size=50).\n\n### Disclaimer:\n**We recommend using [GPUs](https://www.kaggle.com/docs/efficient-gpu-usage) when working with JAX on Kaggle.** These notebooks are compatible with the v3-8 [TPUs](https://www.kaggle.com/docs/tpu) that are provided for free in [Kaggle Notebooks](https://www.kaggle.com/code/new), but JAX was optimized for the newly updated [TPU VM](https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms) architecture which is not yet available on Kaggle.\n","metadata":{}},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"code","source":"#Uncomment and Run when only accelerator is TPU\n#%%capture\n#!conda install -y -c conda-forge jax jaxlib flax optax\n#!conda install -y importlib-metadata","metadata":{"execution":{"iopub.status.busy":"2022-02-21T13:56:56.810276Z","iopub.execute_input":"2022-02-21T13:56:56.810624Z","iopub.status.idle":"2022-02-21T14:10:10.883857Z","shell.execute_reply.started":"2022-02-21T13:56:56.810511Z","shell.execute_reply":"2022-02-21T14:10:10.882929Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Importing all the libraries necessary for the project\nimport os\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\n%matplotlib inline\nimport seaborn as sns\nimport cv2\nimport time\nimport random\nfrom random import randint\nimport time\nimport torch\nfrom torch.utils.data import Dataset, random_split, DataLoader\nimport torch.nn.functional as F\nimport torch.nn as nn\nfrom PIL import Image\nfrom scipy import ndimage\nimport torchvision\nfrom tqdm.notebook import tqdm\nimport tensorflow as tf\nfrom torchvision import transforms\nfrom tqdm import tqdm\nfrom flax.training import train_state\nfrom typing import Any\nimport jax.numpy as jnp\nimport jax.random\nimport functools\nimport optax\nimport flax.linen as nn\nimport jax.nn\n# to suppress warnings caused by cuda version\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:20.462559Z","iopub.execute_input":"2022-02-28T12:06:20.462913Z","iopub.status.idle":"2022-02-28T12:06:28.01885Z","shell.execute_reply.started":"2022-02-28T12:06:20.462816Z","shell.execute_reply":"2022-02-28T12:06:28.018078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### TPU detection and configuration\n**We recommend using [GPUs](https://www.kaggle.com/docs/efficient-gpu-usage) when working with JAX on Kaggle.** These notebooks are compatible with the v3-8 [TPUs](https://www.kaggle.com/docs/tpu) that are provided for free in [Kaggle Notebooks](https://www.kaggle.com/code/new), but JAX was optimized for the newly updated [TPU VM](https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms) architecture which is not yet available on Kaggle.\n","metadata":{"execution":{"iopub.status.busy":"2022-02-04T10:36:27.660527Z","iopub.execute_input":"2022-02-04T10:36:27.660878Z","iopub.status.idle":"2022-02-04T10:38:24.234155Z","shell.execute_reply.started":"2022-02-04T10:36:27.660852Z","shell.execute_reply":"2022-02-04T10:38:24.232535Z"}}},{"cell_type":"code","source":"if '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('Registered TPU:', config.FLAGS.jax_backend_target)\nelse:\n    print('No TPU detected.')","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:34.688659Z","iopub.execute_input":"2022-02-28T12:06:34.688927Z","iopub.status.idle":"2022-02-28T12:06:34.696559Z","shell.execute_reply.started":"2022-02-28T12:06:34.688897Z","shell.execute_reply":"2022-02-28T12:06:34.695592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Shows the list of the available devices\njax.local_devices()","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:36.036658Z","iopub.execute_input":"2022-02-28T12:06:36.036935Z","iopub.status.idle":"2022-02-28T12:06:36.231643Z","shell.execute_reply.started":"2022-02-28T12:06:36.036903Z","shell.execute_reply":"2022-02-28T12:06:36.230885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Reading and Processing\nWe are using Image dataset from famous Kaggle competition [Dog Breed Identification](https://www.kaggle.com/c/dog-breed-identification). This dataset comprises 120 breeds of dogs.\n","metadata":{}},{"cell_type":"code","source":"DATA_DIR = '../input/dog-breed-identification'\nTRAIN_DIR = DATA_DIR + '/train'                           \nTRAIN_CSV = DATA_DIR + '/labels.csv'     \ndata_df = pd.read_csv(TRAIN_CSV)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:38.220062Z","iopub.execute_input":"2022-02-28T12:06:38.220348Z","iopub.status.idle":"2022-02-28T12:06:38.247225Z","shell.execute_reply.started":"2022-02-28T12:06:38.220318Z","shell.execute_reply":"2022-02-28T12:06:38.246515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"labels_names=data_df[\"breed\"].unique()\nlabels_sorted=labels_names.sort()\nlabels = dict(zip(range(len(labels_names)),labels_names))","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:38.976431Z","iopub.execute_input":"2022-02-28T12:06:38.976704Z","iopub.status.idle":"2022-02-28T12:06:38.990253Z","shell.execute_reply.started":"2022-02-28T12:06:38.976673Z","shell.execute_reply":"2022-02-28T12:06:38.989405Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"lbl=[]\npath_img=[]\n\nfor i in range(len(data_df[\"breed\"])):\n    temp1=list(labels.values()).index(data_df.breed[i])\n    lbl.append(temp1)\n    temp2=TRAIN_DIR + \"/\" + str(data_df.id[i]) + \".jpg\"\n    path_img.append(temp2)\n\ndata_df['path_img'] =path_img  \ndata_df['lbl'] = lbl\n\ndata_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:39.722135Z","iopub.execute_input":"2022-02-28T12:06:39.722407Z","iopub.status.idle":"2022-02-28T12:06:39.991643Z","shell.execute_reply.started":"2022-02-28T12:06:39.722376Z","shell.execute_reply":"2022-02-28T12:06:39.990898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Creating custom dataset class using Pytorch's [Datasets and Dataloaders](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#loading-a-dataset)","metadata":{}},{"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, df, root_dir, transform=None):\n        self.df = df\n        self.transform = transform\n        self.root_dir = root_dir\n        \n    def __len__(self):\n        return len(self.df)    \n    \n    def __getitem__(self, idx):\n        row = self.df.loc[idx]\n        img_id, img_label = row['id'], row['lbl']\n        img_fname = self.root_dir + \"/\" + str(img_id) + \".jpg\"\n        img = Image.open(img_fname)\n        if self.transform:\n            img = self.transform(img)\n        return img, img_label","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:40.900121Z","iopub.execute_input":"2022-02-28T12:06:40.900946Z","iopub.status.idle":"2022-02-28T12:06:40.908225Z","shell.execute_reply.started":"2022-02-28T12:06:40.900897Z","shell.execute_reply":"2022-02-28T12:06:40.907334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"As we proceed, we will be applying data augmentations to the dataset using Pytorch's [transforms](https://pytorch.org/vision/stable/transforms.html) and returning data in the form of numpy arrays","metadata":{}},{"cell_type":"code","source":"IMAGE_HEIGHT = 128\nIMAGE_WIDTH = 128\n\ntraining_transform = transforms.Compose([\n    transforms.RandomAffine(degrees=(-30, 30),\n                            translate=(0.0, 0.2)),\n    transforms.RandomHorizontalFlip(),\n    transforms.Resize((IMAGE_HEIGHT,\n                        IMAGE_WIDTH)),\n    np.array])\n\ntesting_transform = transforms.Compose([\n    transforms.Resize((IMAGE_HEIGHT,\n                       IMAGE_WIDTH)),\n    np.array])\n\n\nnp.random.seed(42)\nmsk = np.random.rand(len(data_df)) < 0.8\n\ntrain_df = data_df[msk].reset_index()\nval_df = data_df[~msk].reset_index()\n\ntrain_ds = DogDataset(train_df, TRAIN_DIR, transform=training_transform)\nval_ds = DogDataset(val_df, TRAIN_DIR, transform=testing_transform)\nlen(train_ds), len(val_ds)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:41.90328Z","iopub.execute_input":"2022-02-28T12:06:41.903547Z","iopub.status.idle":"2022-02-28T12:06:41.921671Z","shell.execute_reply.started":"2022-02-28T12:06:41.903516Z","shell.execute_reply":"2022-02-28T12:06:41.920892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The Dataset retrieves our dataset’s features and labels one sample at a time. While training a model, we typically want to pass samples in “minibatches”, reshuffle the data at every epoch to reduce model overfitting, and use Python’s multiprocessing to speed up data retrieval.\n\nDataLoader is an iterable that abstracts this complexity for us in an easy API.\n\n-Source: https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#loading-a-dataset","metadata":{}},{"cell_type":"code","source":"import jax\nNUM_TPUS = jax.device_count()\nBATCH_SIZE = 128\ntrain_dataloader = DataLoader(train_ds,\n                                               batch_size=BATCH_SIZE,\n                                               shuffle=True, drop_last=True,\n                                               num_workers=0)\ntest_dataloader = DataLoader(val_ds,\n                                              batch_size=BATCH_SIZE,\n                                              shuffle=True, drop_last=True,\n                                              num_workers=0)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:43.034327Z","iopub.execute_input":"2022-02-28T12:06:43.035033Z","iopub.status.idle":"2022-02-28T12:06:43.040804Z","shell.execute_reply.started":"2022-02-28T12:06:43.034981Z","shell.execute_reply":"2022-02-28T12:06:43.0401Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"(image_batch, label_batch) = next(iter(train_dataloader))\nprint(image_batch.shape)\nprint(label_batch.shape)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:44.050918Z","iopub.execute_input":"2022-02-28T12:06:44.051336Z","iopub.status.idle":"2022-02-28T12:06:45.63049Z","shell.execute_reply.started":"2022-02-28T12:06:44.051297Z","shell.execute_reply":"2022-02-28T12:06:45.629641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating batches/shrades of the data\n\nReading batches of data from the CPU's RAM and copying it to the memory of the accelerator which you're going to use for computation in the form of ShardedDeviceArray(s) using JAX's [`device_put_sharded`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put_sharded.html#jax.device_put_sharded).","metadata":{}},{"cell_type":"code","source":"def copy_dataset_to_devices(dataset, devices, num_reps=1):\n    sharded_images = []\n    sharded_labels = []\n    for _ in range(num_reps):\n        for image_batch, label_batch in tqdm(dataset, ncols=100):\n            image_batch = image_batch.detach().cpu().numpy()\n            image_batches = np.split(image_batch, NUM_TPUS, axis = 0)\n            sharded_device_images = jax.device_put_sharded(image_batches, devices)\n            sharded_images.append(sharded_device_images)\n\n            label_batch = label_batch.detach().cpu().numpy()\n            label_batches = np.split(label_batch, NUM_TPUS, axis = 0)\n            sharded_device_labels = jax.device_put_sharded(label_batches, devices)\n            sharded_labels.append(sharded_device_labels)\n\n    return sharded_images, sharded_labels\n\ndevices = jax.local_devices()\nsharded_training_images, sharded_training_labels = copy_dataset_to_devices(train_dataloader, devices, num_reps=10)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:06:47.192107Z","iopub.execute_input":"2022-02-28T12:06:47.19238Z","iopub.status.idle":"2022-02-28T12:15:06.634723Z","shell.execute_reply.started":"2022-02-28T12:06:47.19235Z","shell.execute_reply":"2022-02-28T12:15:06.634007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model architecture\n\nHere, in this notebook we'll be using VGG19 network. we'll be using [FLAX Linen package](https://flax.readthedocs.io/en/latest/flax.linen.html) for defining the model architecture from scratch.","metadata":{}},{"cell_type":"code","source":"NUM_CLASSES = 120 \nclass VGG19(nn.Module):\n    @nn.compact\n    def __call__(self, x, training):\n        x = self._stack(x, 64, training)\n        x = self._stack(x, 64, training)\n        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))\n    \n        x = self._stack(x, 128, training)\n        x = self._stack(x, 128, training)\n        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))\n\n        x = self._stack(x, 256, training)\n        x = self._stack(x, 256, training)\n        x = self._stack(x, 256, training)\n        x = self._stack(x, 256, training)\n        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))    \n\n        x = self._stack(x, 512, training)\n        x = self._stack(x, 512, training)\n        x = self._stack(x, 512, training)\n        x = self._stack(x, 512, training)\n        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))    \n    \n        x = self._stack(x, 512, training)\n        x = self._stack(x, 512, training)\n        x = self._stack(x, 512, training)\n        x = self._stack(x, 512, training)\n        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))  \n\n        x = x.reshape((x.shape[0], -1))\n\n        x = nn.Dense(features=4096)(x)\n        x = nn.BatchNorm(use_running_average=not training)(x)\n        x = nn.relu(x)\n        x = nn.Dropout(0.5, deterministic=not training)(x)\n\n        x = nn.Dense(features=4096)(x)\n        x = nn.BatchNorm(use_running_average=not training)(x)\n        x = nn.relu(x)\n        x = nn.Dropout(0.5, deterministic=not training)(x)\n    \n        x = nn.Dense(features=NUM_CLASSES)(x)\n        x = nn.log_softmax(x)\n        return x\n  \n    @staticmethod\n    def _stack(x, features, training, dropout=None):\n        x = nn.Conv(features=features, kernel_size=(3, 3), padding='SAME')(x)\n        x = nn.BatchNorm(use_running_average=not training)(x)\n        x = nn.relu(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:10.027421Z","iopub.execute_input":"2022-02-28T12:15:10.027898Z","iopub.status.idle":"2022-02-28T12:15:10.046672Z","shell.execute_reply.started":"2022-02-28T12:15:10.027856Z","shell.execute_reply":"2022-02-28T12:15:10.044018Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train function\nIn the train function, we'll collect a batch of train data by looping through sharded training images and sharded training labels to train our neural network from the given state, and we'll get back our new training state as well as batch statistics.","metadata":{}},{"cell_type":"code","source":"def average_metrics(metrics):\n    '''\n    Takes the list of dictionaries of the form k: v, and returns a dictionary\n     of the form k: (average of the v).\n    '''\n    return {k: np.mean([metric[k] for metric in metrics])\n        for k in metrics[0]}\n\ndef train(initial_network_state, num_epochs):\n    '''\n    Training the model from the given state, returns the state along with the training accuracies\n    '''\n    training_accuracies = []\n    state = initial_network_state\n    for i in range(num_epochs):\n        batch_metrics = []\n        for (image_batch, label_batch) in tqdm(zip(sharded_training_images,\n                                               sharded_training_labels),\n                                           total=len(sharded_training_images),\n                                           ncols=100):\n            state, metrics = train_batch(state, image_batch, label_batch)\n            batch_metrics.append(metrics)\n        train_metrics = average_metrics(batch_metrics)\n        print(f'Epoch {i+1} done.', flush=True)\n        print(f'  Loss: {train_metrics[\"loss\"]:.4f}, '\n          + f'accuracy: {train_metrics[\"accuracy\"]:.4f}', flush=True)\n        training_accuracies.append(train_metrics[\"accuracy\"])\n    return state, training_accuracies","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:11.294553Z","iopub.execute_input":"2022-02-28T12:15:11.29508Z","iopub.status.idle":"2022-02-28T12:15:11.303401Z","shell.execute_reply.started":"2022-02-28T12:15:11.295041Z","shell.execute_reply":"2022-02-28T12:15:11.302546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model Initialization functions\n\nIn FLAX, we have to manually create and update `train_state` which holds all the model's variables using [`flax.training.train_state`](https://flax.readthedocs.io/en/latest/_modules/flax/training/train_state.html#TrainState)","metadata":{}},{"cell_type":"code","source":"class VGGState(train_state.TrainState):\n    rng: Any\n    batch_stats: Any\n  \n    @classmethod\n    def create(cls, apply_fn, params, tx, rng, batch_stats):\n        opt_state = tx.init(params)\n        state = cls(0, apply_fn, params, tx, opt_state, rng, batch_stats)\n        return state\n  \n    @classmethod\n    def update_rng(cls, state, rng):\n        return VGGState.create(state.apply_fn, state.params, state.tx, rng,\n                           state.batch_stats)\n  \n    @classmethod\n    def update_batch_stats(cls, state, batch_stats):\n        return VGGState.create(state.apply_fn, state.params, state.tx,\n                           state.rng, batch_stats)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:13.029463Z","iopub.execute_input":"2022-02-28T12:15:13.030266Z","iopub.status.idle":"2022-02-28T12:15:13.0391Z","shell.execute_reply.started":"2022-02-28T12:15:13.030212Z","shell.execute_reply":"2022-02-28T12:15:13.038353Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss & Metrics calculations\nNow, we will define the functions which calculates the training loss and the accuracy using the given predicted values and labels","metadata":{}},{"cell_type":"code","source":"def accuracy(logits, labels):\n    '''\n    Calcualtes the accuracy using the given logits and labels\n    '''\n    return jnp.mean(jnp.argmax(logits, -1) == labels)\n\ndef cross_entropy(logits, labels):\n    '''\n    Cross Entropy error between the logits and labels\n    '''\n    one_hot_labels = jax.nn.one_hot(labels, NUM_CLASSES)\n    cross_entropy = optax.softmax_cross_entropy(logits, one_hot_labels)\n    return jnp.mean(cross_entropy)\n\ndef training_loss(image_batch, label_batch, rng, batch_stats, params):\n    '''\n    Calculates the training loss \n    '''\n    logits, batch_stats = VGG19().apply({'params': params,\n                                       'batch_stats': batch_stats},\n                                      image_batch, \n                                      training=True,\n                                      rngs={'dropout': rng},\n                                      mutable=['batch_stats'])\n    loss = cross_entropy(logits, label_batch)\n    return loss, (logits, batch_stats)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:14.914343Z","iopub.execute_input":"2022-02-28T12:15:14.914608Z","iopub.status.idle":"2022-02-28T12:15:14.923466Z","shell.execute_reply.started":"2022-02-28T12:15:14.914577Z","shell.execute_reply":"2022-02-28T12:15:14.922625Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training a single batch function\nWe will now define the function for training a single batch of data, which will take the current train state and the training data as input and return the updated train state along with the training statistics.","metadata":{}},{"cell_type":"code","source":"@functools.partial(jax.pmap, axis_name='tpu')\ndef train_batch(state, image_batch, label_batch):\n    rng, subrng = jax.random.split(state.rng)\n    batch_loss_fn = functools.partial(training_loss, image_batch, label_batch,\n                                    subrng, state.batch_stats)\n    (batch_loss, (logits, batch_stats)), grads = \\\n    jax.value_and_grad(batch_loss_fn, has_aux=True)(state.params)\n  \n    gradsum = jax.lax.psum(grads, axis_name='tpu')\n    \n    state = state.apply_gradients(grads=gradsum)\n    state = state.update_batch_stats(state, batch_stats['batch_stats'])\n    state = state.update_rng(state, rng)\n\n    batch_accuracy = accuracy(logits=logits, labels=label_batch)\n    batch_accuracy_sum = jax.lax.pmean(batch_accuracy, axis_name='tpu')\n    batch_loss = jax.lax.psum(batch_loss, axis_name='tpu')\n    stats = {'loss': batch_loss, 'accuracy': batch_accuracy_sum}  \n    return state, stats","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:16.800366Z","iopub.execute_input":"2022-02-28T12:15:16.801002Z","iopub.status.idle":"2022-02-28T12:15:16.808922Z","shell.execute_reply.started":"2022-02-28T12:15:16.800958Z","shell.execute_reply":"2022-02-28T12:15:16.808132Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating train state\nCreating the initial train state which we'll be passing to the neural network while training","metadata":{}},{"cell_type":"code","source":"def create_train_state(rng, dummy_image_batch):\n    net = VGG19()\n    params = net.init({'params': rng, 'dropout': rng}, dummy_image_batch, True)\n    tx = optax.adam(learning_rate=0.001)\n    state = VGGState.create(net.apply, params['params'], tx, rng,\n                          params['batch_stats'])\n    return state","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:20.187255Z","iopub.execute_input":"2022-02-28T12:15:20.187519Z","iopub.status.idle":"2022-02-28T12:15:20.193033Z","shell.execute_reply.started":"2022-02-28T12:15:20.187489Z","shell.execute_reply":"2022-02-28T12:15:20.191964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rng = jax.random.PRNGKey(42)\nrngs = np.broadcast_to(rng, (NUM_TPUS,) + rng.shape)\nsome_dummy_image_batch = sharded_training_images[0]\nstate = jax.pmap(create_train_state, axis_name='tpu')(rngs,some_dummy_image_batch)","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:22.710445Z","iopub.execute_input":"2022-02-28T12:15:22.710709Z","iopub.status.idle":"2022-02-28T12:15:32.989595Z","shell.execute_reply.started":"2022-02-28T12:15:22.710678Z","shell.execute_reply":"2022-02-28T12:15:32.988855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training \nNext, we will train the VGG19 neural network for 75 epochs and plot the accuracy graph to see how well the model does.","metadata":{}},{"cell_type":"code","source":"start = time.time()\nfinal_state, training_accuracies = train(state, num_epochs=75)\nprint(\"Total time: \", time.time() - start, \"seconds\")","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:15:46.634824Z","iopub.execute_input":"2022-02-28T12:15:46.635608Z","iopub.status.idle":"2022-02-28T12:16:22.446888Z","shell.execute_reply.started":"2022-02-28T12:15:46.635567Z","shell.execute_reply":"2022-02-28T12:16:22.445797Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot the Accuracy \nplt.plot(training_accuracies)\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Accuracy\")\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-02-28T12:18:31.383519Z","iopub.execute_input":"2022-02-28T12:18:31.384376Z","iopub.status.idle":"2022-02-28T12:18:31.410214Z","shell.execute_reply.started":"2022-02-28T12:18:31.38432Z","shell.execute_reply":"2022-02-28T12:18:31.408717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **Conclusion**\nHere in this notebook, we've illustrated how [JAX](https://github.com/google/jax) and [FLAX](https://flax.readthedocs.io/en/latest/) can be used to train the neural network from scratch for the image classification dataset, with an accuracy of more than 80%. To see more examples of how to use [JAX](https://github.com/google/jax) and [FLAX](https://flax.readthedocs.io/en/latest/) with different data formats, please see this discussion post.  \n\nNow, it's your turn to  create some amazing notebooks using [JAX](https://github.com/google/jax) and [FLAX](https://flax.readthedocs.io/en/latest/). \n\n### **Useful resources which helped me:**\n\n* https://flax.readthedocs.io/en/latest/notebooks/annotated_mnist.html\n* https://jax.readthedocs.io/en/latest/\n* https://gist.github.com/fedelebron/b7be87a4feb88786cc142ef99931ff06#file-dog-classifier-ipynb\n* https://www.kaggle.com/anujverma126/dog-breed-classification-beginner-s-tutorial \n* https://github.com/google/flax/tree/main/examples/imagenet","metadata":{}}]}