{"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-02-16T16:13:06.260802Z","iopub.execute_input":"2022-02-16T16:13:06.261115Z","iopub.status.idle":"2022-02-16T16:14:27.698864Z","shell.execute_reply.started":"2022-02-16T16:13:06.261072Z","shell.execute_reply":"2022-02-16T16:14:27.698013Z"}}},{"cell_type":"markdown","source":"## Imports\nUncomment and Run this code cell when only accelerator is TPU","metadata":{}},{"cell_type":"code","source":"#%%capture\n#!conda install -y -c conda-forge jax jaxlib flax optax datasets transformers\n#!conda install -y importlib-metadata","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:18:02.233607Z","iopub.execute_input":"2022-02-24T07:18:02.234438Z","iopub.status.idle":"2022-02-24T07:22:27.211456Z","shell.execute_reply.started":"2022-02-24T07:18:02.2343Z","shell.execute_reply":"2022-02-24T07:22:27.209666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Importing all the libraries necessary for the project\nimport os\nimport sys\nimport time\nimport math\nimport random\nimport librosa\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport cv2\nimport matplotlib.pyplot as plt\n%matplotlib inline\nimport seaborn as sns\nimport torch\nimport torch.nn as nn\nimport torch.nn.init as init\nfrom torchvision import transforms, utils\nfrom torch.utils.data import Dataset, DataLoader\n# Ignore warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport tensorflow as tf\nfrom torch.utils.data import Dataset, random_split, DataLoader\nimport gc\nfrom tqdm import tqdm, tqdm_notebook; tqdm.pandas()\nimport jax\nimport torchvision\nimport optax\nimport flax.linen as nn\nimport jax.nn\nimport jax.numpy as jnp\nfrom flax import linen as nn\nfrom tensorflow.keras.utils import to_categorical\nseed = 1234\nnp.random.seed(seed)\nimport warnings\ndef fxn():\n    warnings.warn(\"deprecated\", DeprecationWarning)\nwith warnings.catch_warnings():\n    warnings.simplefilter(\"ignore\")\n    fxn()\nwarnings.filterwarnings(\"ignore\")\nfrom torchvision import transforms\nimport torch\nfrom typing import Any\nimport functools\nfrom flax.training import train_state\n# to suppress warnings caused by cuda version\nos.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:23:36.352199Z","iopub.execute_input":"2022-02-24T07:23:36.352947Z","iopub.status.idle":"2022-02-24T07:23:36.41053Z","shell.execute_reply.started":"2022-02-24T07:23:36.35289Z","shell.execute_reply":"2022-02-24T07:23:36.409205Z"},"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":{}},{"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\nexcept ValueError:\n    tpu = None\n    \nif tpu:\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-24T07:22:37.028849Z","iopub.status.idle":"2022-02-24T07:22:37.029525Z","shell.execute_reply.started":"2022-02-24T07:22:37.029302Z","shell.execute_reply":"2022-02-24T07:22:37.029325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# List all the available devices\njax.devices()","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.031047Z","iopub.status.idle":"2022-02-24T07:22:37.031386Z","shell.execute_reply.started":"2022-02-24T07:22:37.031214Z","shell.execute_reply":"2022-02-24T07:22:37.031235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load and Pre-process the dataset\nFor time and memory management, we'll be taking random sample of 15 birds","metadata":{}},{"cell_type":"code","source":"# seeding function for reproducibility\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.034504Z","iopub.status.idle":"2022-02-24T07:22:37.034835Z","shell.execute_reply.started":"2022-02-24T07:22:37.034676Z","shell.execute_reply":"2022-02-24T07:22:37.034691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ROOT = \"/kaggle/input/birdsong-recognition/\"\nos.listdir(ROOT)","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.036111Z","iopub.status.idle":"2022-02-24T07:22:37.036416Z","shell.execute_reply.started":"2022-02-24T07:22:37.036257Z","shell.execute_reply":"2022-02-24T07:22:37.036272Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv(os.path.join(ROOT, 'train.csv'))[['ebird_code', 'filename', 'duration']]\ndf['path'] = ROOT+'train_audio/' + df['ebird_code'] + \"/\" + df['filename']\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.037586Z","iopub.status.idle":"2022-02-24T07:22:37.037888Z","shell.execute_reply.started":"2022-02-24T07:22:37.037729Z","shell.execute_reply":"2022-02-24T07:22:37.037743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SEED = 42\nFRAC = 0.2     # Validation fraction\nSR = 44100     # sampling rate\nMAXLEN= 60    # seconds\nN_MELS = 128\n\nseed_everything(SEED)\ndevice = torch.device('cpu')\n\n#Random sample of 15 birds\nclasses = set(random.sample(df['ebird_code'].unique().tolist(), 15)) \nprint(classes)","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.039Z","iopub.status.idle":"2022-02-24T07:22:37.039321Z","shell.execute_reply.started":"2022-02-24T07:22:37.039144Z","shell.execute_reply":"2022-02-24T07:22:37.039173Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = df[df.ebird_code.apply(lambda x: x in classes)].reset_index(drop=True)\nkeys = set(df.ebird_code)\nvalues = np.arange(0, len(keys))\ncode_dict = dict(zip(sorted(keys), values))\ndf['label'] = df['ebird_code'].apply(lambda x: code_dict[x])\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.040211Z","iopub.status.idle":"2022-02-24T07:22:37.040509Z","shell.execute_reply.started":"2022-02-24T07:22:37.040355Z","shell.execute_reply":"2022-02-24T07:22:37.04037Z"},"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 BirdSoundDataset(Dataset):\n    \"\"\"Bird Sound dataset.\"\"\"\n\n    def __init__(self, df, transform = None):\n        \"\"\"\n        Args:\n            df (pd.DataFrame): must have ['path', 'label'] columns\n        \"\"\"\n        self.df = df\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df)\n    \n    \n    def loadMP3(self, path, duration):\n        \"\"\"\n        Args:\n            path: path of the audio file \n        Returns:\n            mels: Melspectrogram of the given audio file \n        \"\"\"\n        try:\n            duration=5\n            samples = SR* duration\n            audio, _ = librosa.load(path, sr=SR)\n            \n            if 0 < len(audio):\n                audio, _ = librosa.effects.trim(audio)\n            if len(audio) > samples: # long enough\n                audio = audio[0:0+samples]\n            else: # pad blank\n                padding = samples - len(audio)\n                offset = padding // 2\n                y = np.pad(audio, (offset, samples - len(audio) - offset), 'constant')\n\n            mels = librosa.feature.melspectrogram(y=audio, sr=SR,n_mels=N_MELS, hop_length = 347,n_fft = N_MELS *20,fmin = 20, fmax = SR//2)\n            mels = librosa.power_to_db(mels).astype(np.float32)\n            mels = mels.transpose()\n            eps = 0.001\n            if np.std(mels) != 0:\n                mels = (mels - np.mean(mels)) / np.std(mels)\n            else:\n                mels = (mels - np.mean(mels)) / eps\n            return mels\n            \n        except Exception as e:\n            print(\"Error encountered while parsing file: \", path, e)\n            mels = np.zeros((N_MELS, MAXLEN*SR//347), dtype=np.float32)\n            return mels\n            \n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        \n        path = self.df['path'].iloc[idx]\n    \n        duration=5\n        if os.path.exists(\"./\"+path.split('/')[-1]+\".npy\"):\n            mels = np.load(\"./\"+path.split('/')[-1]+\".npy\")\n        else:\n            \n            mels = self.loadMP3(path, duration)\n            np.save(\"./\"+path.split('/')[-1]+\".npy\", mels)\n        label  = self.df['label'].iloc[idx]\n        mels = np.resize(mels,(636,128,1))\n        return mels, label","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.041776Z","iopub.status.idle":"2022-02-24T07:22:37.042094Z","shell.execute_reply.started":"2022-02-24T07:22:37.041933Z","shell.execute_reply":"2022-02-24T07:22:37.04195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dividing the dataset into train and validation sets\ndf = df.sample(frac=1).reset_index(drop=True)\ntrain_len = int(len(df) * (1-FRAC))\ntrain_df = df.iloc[:train_len]\nvalid_df = df.iloc[train_len:]\ntrain_df.shape, valid_df.shape","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.043654Z","iopub.status.idle":"2022-02-24T07:22:37.044165Z","shell.execute_reply.started":"2022-02-24T07:22:37.043886Z","shell.execute_reply":"2022-02-24T07:22:37.04393Z"},"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":"# prepare data loaders \n#NUM_TPUS = jax.device_count()\nBATCH_SIZE = 32\n\ntrain_loader = torch.utils.data.DataLoader(BirdSoundDataset(train_df),\n                                           batch_size=BATCH_SIZE, \n                                           num_workers=0, \n                                           shuffle=True, \n                                           drop_last = True)\n\nvalid_loader = torch.utils.data.DataLoader(BirdSoundDataset(valid_df), \n                                           batch_size=BATCH_SIZE, \n                                           num_workers=0, \n                                           shuffle=True, \n                                           drop_last = True)\n\nlen(train_loader), len(valid_loader)","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.045509Z","iopub.status.idle":"2022-02-24T07:22:37.046008Z","shell.execute_reply.started":"2022-02-24T07:22:37.045724Z","shell.execute_reply":"2022-02-24T07:22:37.045748Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"(image_batch, label_batch) = next(iter(train_loader))\nprint(image_batch.shape)\nprint(label_batch.shape)","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.04764Z","iopub.status.idle":"2022-02-24T07:22:37.048135Z","shell.execute_reply.started":"2022-02-24T07:22:37.047868Z","shell.execute_reply":"2022-02-24T07:22:37.04789Z"},"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":"NUM_TPUS = jax.device_count()\n\ndef 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_loader, devices, num_reps=10)","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.049449Z","iopub.status.idle":"2022-02-24T07:22:37.049937Z","shell.execute_reply.started":"2022-02-24T07:22:37.049659Z","shell.execute_reply":"2022-02-24T07:22:37.049681Z"},"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 = 15 \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 = 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        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-24T07:22:37.051605Z","iopub.status.idle":"2022-02-24T07:22:37.052117Z","shell.execute_reply.started":"2022-02-24T07:22:37.051822Z","shell.execute_reply":"2022-02-24T07:22:37.051846Z"},"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-24T07:22:37.053374Z","iopub.status.idle":"2022-02-24T07:22:37.054021Z","shell.execute_reply.started":"2022-02-24T07:22:37.053721Z","shell.execute_reply":"2022-02-24T07:22:37.053748Z"},"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-24T07:22:37.055794Z","iopub.status.idle":"2022-02-24T07:22:37.056371Z","shell.execute_reply.started":"2022-02-24T07:22:37.056091Z","shell.execute_reply":"2022-02-24T07:22:37.056118Z"},"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-24T07:22:37.058207Z","iopub.status.idle":"2022-02-24T07:22:37.058704Z","shell.execute_reply.started":"2022-02-24T07:22:37.058456Z","shell.execute_reply":"2022-02-24T07:22:37.05848Z"},"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    '''\n    Training a single batch and returns loss and the accuracy\n    '''\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\n    return state, stats","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.060023Z","iopub.status.idle":"2022-02-24T07:22:37.060333Z","shell.execute_reply.started":"2022-02-24T07:22:37.060173Z","shell.execute_reply":"2022-02-24T07:22:37.060189Z"},"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.01)\n    state = VGGState.create(net.apply, params['params'], tx, rng,\n                          params['batch_stats'])\n    return state","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.061194Z","iopub.status.idle":"2022-02-24T07:22:37.061478Z","shell.execute_reply.started":"2022-02-24T07:22:37.061322Z","shell.execute_reply":"2022-02-24T07:22:37.061336Z"},"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-24T07:22:37.062224Z","iopub.status.idle":"2022-02-24T07:22:37.062504Z","shell.execute_reply.started":"2022-02-24T07:22:37.062352Z","shell.execute_reply":"2022-02-24T07:22:37.062366Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training \nNext, we will train the VGG19 neural network for 25 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=25)\nprint(\"Total time: \", time.time() - start, \"seconds\")","metadata":{"execution":{"iopub.status.busy":"2022-02-24T07:22:37.063246Z","iopub.status.idle":"2022-02-24T07:22:37.063753Z","shell.execute_reply.started":"2022-02-24T07:22:37.063536Z","shell.execute_reply":"2022-02-24T07:22:37.06356Z"},"_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-24T07:22:37.064851Z","iopub.status.idle":"2022-02-24T07:22:37.065181Z","shell.execute_reply.started":"2022-02-24T07:22:37.065021Z","shell.execute_reply":"2022-02-24T07:22:37.065036Z"},"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 audio classification dataset, with an accuracy of more than 95%. 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* https://www.kaggle.com/nilaychauhan/convert-cornell-birdcall-recognition-to-tfrecords\n* https://www.kaggle.com/servietsky/fast-import-audio-and-save-spectrograms/notebook\n* https://www.kaggle.com/dhananjay3/simple-pytorch-starter/notebook\n* https://github.com/google/flax/tree/main/examples/imagenet\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","metadata":{}}]}