flax_mnist.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "flax_mnist.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMAMrSyi2SMsFF6tOZIv9Kf",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deterministic-algorithms-lab/Jax-Journey/blob/main/flax_mnist.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "612x-KQNev6L",
        "outputId": "f45e417c-ed3c-4af1-82be-232e6af172c1"
      },
      "source": [
        "# Install ml-collections & latest Flax version from Github.\n",
        "!pip install -q ml-collections git+https://github.com/google/flax"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[?25l\r\u001b[K     |███▊                            | 10kB 14.9MB/s eta 0:00:01\r\u001b[K     |███████▍                        | 20kB 19.8MB/s eta 0:00:01\r\u001b[K     |███████████                     | 30kB 22.6MB/s eta 0:00:01\r\u001b[K     |██████████████▉                 | 40kB 16.8MB/s eta 0:00:01\r\u001b[K     |██████████████████▌             | 51kB 10.7MB/s eta 0:00:01\r\u001b[K     |██████████████████████▏         | 61kB 12.0MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▉      | 71kB 10.9MB/s eta 0:00:01\r\u001b[K     |█████████████████████████████▋  | 81kB 11.9MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 92kB 6.3MB/s \n",
            "\u001b[?25h  Building wheel for flax (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DaPx_QXfe_q4"
      },
      "source": [
        "ML Collections is a library of collections(like normal python ```collections``` module) specialised for ML. The repo can be viewed [here](https://github.com/google/ml_collections)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "96r859yCeb5l"
      },
      "source": [
        "import ml_collections\n",
        "\n",
        "def get_config():\n",
        "    config = ml_collections.ConfigDict()\n",
        "    \n",
        "    config.learning_rate = 0.1\n",
        "    config.momentum = 0.9\n",
        "    config.batch_size = 128\n",
        "    config.num_epochs = 10\n",
        "\n",
        "    return config"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4RbQ2kMRnP5e"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EaDe0Bimff15"
      },
      "source": [
        "from absl import logging\n",
        "import flax\n",
        "import jax.numpy as jnp\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "logging.set_verbosity(logging.INFO)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0pdBQBGhfy-p"
      },
      "source": [
        "# Helper functions for images.\n",
        "\n",
        "def show_img(img, ax=None, title=None):\n",
        "  \"\"\"Shows a single image.\"\"\"\n",
        "  if ax is None:\n",
        "    ax = plt.gca()\n",
        "  ax.imshow(img[..., 0], cmap='gray')\n",
        "  ax.set_xticks([])\n",
        "  ax.set_yticks([])\n",
        "  if title:\n",
        "    ax.set_title(title)\n",
        "\n",
        "def show_img_grid(imgs, titles):\n",
        "  \"\"\"Shows a grid of images.\"\"\"\n",
        "  n = int(np.ceil(len(imgs)**.5))\n",
        "  _, axs = plt.subplots(n, n, figsize=(3 * n, 3 * n))\n",
        "  for i, (img, title) in enumerate(zip(imgs, titles)):\n",
        "    show_img(img, axs[i // n][i % n], title)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5wo9zShEf6ts"
      },
      "source": [
        "# Local imports from current directory will auto reload.\n",
        "# Any changes you make to local files will appear automatically.\n",
        "%load_ext autoreload\n",
        "%autoreload 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3iwWpea0gGdZ"
      },
      "source": [
        "config = get_config()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eaed4ugFiLNx"
      },
      "source": [
        "* ```tfds.as_numpy()``` takes in a dataset to a python generator, that generates numpy matrices here. \n",
        "\n",
        "* ```tf.DatasetBuilder.as_dataset()``` builds an input pipeline(taking care of all batch size, device etc.) using ```tf.data.Dataset```(s). The ```tf.data.Dataset```(s) correspond to the ```nn.Dataset``` of PyTorch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nDNet9fwgOEW"
      },
      "source": [
        "def get_datasets():\n",
        "    ds_builder = tfds.builder('mnist')\n",
        "    ds_builder.download_and_prepare()\n",
        "    train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))\n",
        "    test_ds = tfds.as_numpy(ds_builder.as_dataset(split='test', batch_size=-1))\n",
        "    #print(test_ds)                                                             #Each dataset has data in different format, so do check.Here the structure is a dict {'image':np array of all images, 'label': all labels}\n",
        "    #print(train_ds['image'][0])                                                #Prints first image. The values are 0 to 255..\n",
        "    train_ds['image'] = jnp.float32(train_ds['image']) /255 \n",
        "    test_ds['image'] = jnp.float32(test_ds['image']) / 255\n",
        "    return train_ds, test_ds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ONPZKDz0m5aW"
      },
      "source": [
        "train_ds, test_ds = get_datasets()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QgJcjvfqm8y4"
      },
      "source": [
        "show_img_grid(\n",
        "    [train_ds['image'][idx] for idx in range(25)],\n",
        "    [f'label={train_ds[\"label\"][idx]}' for idx in range(25)],\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cMKtqEuAov-s"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "auCejTm_ql1C"
      },
      "source": [
        "from flax import linen as nn\n",
        "from flax import optim\n",
        "from flax.metrics import tensorboard\n",
        "import numpy as onp\n",
        "from jax import random\n",
        "import jax"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wIqCGxm5oz4d"
      },
      "source": [
        "class CNN(nn.Module):\n",
        "    @nn.compact\n",
        "    def __call__(self, x):\n",
        "        x = nn.Conv(features=32, kernel_size=(3,3))(x)\n",
        "        x = nn.relu(x)\n",
        "        x = nn.avg_pool(x, window_shape=(2,2), strides=(2,2))\n",
        "        x = nn.Conv(features=64, kernel_size=(3,3))(x)\n",
        "        x = nn.relu(x)\n",
        "        x = nn.avg_pool(x, window_shape=(2,2), strides=(2,3))\n",
        "        x = x.reshape((x.shape[0], -1))\n",
        "        x = nn.Dense(features=256)(x)\n",
        "        x = nn.relu(x)\n",
        "        x = nn.Dense(features=10)(x)\n",
        "        x = nn.log_softmax(x)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "17NZsG4pqzPg"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "key1, key2 = random.split(key)\n",
        "x = random.normal(key1, (1, 28, 28, 1))\n",
        "\n",
        "model = CNN()\n",
        "params = model.init(key2, x)\n",
        "print(params)                                                                   #To check dictionary structure.. whether variables are there, etc."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XxlG23AgtxnW"
      },
      "source": [
        "# Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sk6kUQs4tmyo"
      },
      "source": [
        "optimizer_def = optim.Momentum(learning_rate=config.learning_rate, \n",
        "                               beta=config.momentum)\n",
        "optimizer = optimizer_def.create(params)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WQFafZjjnp1s"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kpSDfgJevg1S"
      },
      "source": [
        "##Loss Funtion"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j6tYKrhVwRf2"
      },
      "source": [
        "def cross_entropy_loss(labels,logits):\n",
        "    return -jnp.mean(jnp.sum(labels*logits, axis=-1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Z1ECDuzwvMY"
      },
      "source": [
        "max_classes=10\n",
        "def onehot(label):\n",
        "    x = (label[...,None]==jnp.arange(0,max_classes)[None])\n",
        "    return x.astype(jnp.float32)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_IcIm-CmvgC7"
      },
      "source": [
        "def loss_fn(params, batch):                                                     #Can input any number of arguments.\n",
        "    logits = CNN().apply(params, batch['image'])                                #We are not constrained to use the same model as before.\n",
        "    loss = cross_entropy_loss(onehot(batch['label']), logits)\n",
        "    return loss, logits                                                         #Can output at most two values"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZxSZHTpk858v"
      },
      "source": [
        "## Metric Calculation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SAe5DQt99CaY"
      },
      "source": [
        "def compute_metric(logits, labels):\n",
        "    loss = cross_entropy_loss(logits, onehot(labels))\n",
        "    accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n",
        "    metrics = {\n",
        "        'loss' : loss,\n",
        "        'accuracy' : accuracy,\n",
        "    }\n",
        "    return metrics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i8HgXZ9UvSZK"
      },
      "source": [
        "## Single Step Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xM-h7BX56yI7"
      },
      "source": [
        "The ```has_aux=True``` below is necessary to indicate that the ```loss_fn``` returns two values, first of which is output of mathematical operation and second is auxillary data. The inability to print any the abstractions adopted by ```JAX``` are very nicely explained [here]( https://github.com/google/jax/issues/196#issuecomment-451671635 ) ."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L7s5bBYbvRhb"
      },
      "source": [
        "@jax.jit\n",
        "def train_step(optimizer, batch):\n",
        "    grad_n_val_fn = jax.value_and_grad(loss_fn, has_aux=True)                   #By default, gradients will be calculated w.r.t the first argument of loss_fn only. \n",
        "    (loss, logits), grad = grad_n_val_fn(optimizer.target, batch)\n",
        "    optimizer = optimizer.apply_gradient(grad)\n",
        "    \n",
        "    #print(loss)                                                                #Not able to get value of loss directly. \n",
        "                                                                                #Can't print values inside jit compiled functions and others nested,inside it, yet.\n",
        "    return optimizer, compute_metric(logits, batch['label'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xe-MMzbIz5Nc"
      },
      "source": [
        "## Epoch Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ML4mqrmi0skB"
      },
      "source": [
        "### Setting up data loading"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G-_Yht9Qz4yQ"
      },
      "source": [
        "train_ds_size = len(train_ds['image'])\n",
        "steps_per_epoch = train_ds_size//config.batch_size\n",
        "\n",
        "perms = random.permutation(key, len(train_ds['image']))\n",
        "perms = perms[:steps_per_epoch*config.batch_size]\n",
        "perms = perms.reshape((steps_per_epoch, config.batch_size))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "na-jvDS50xmW"
      },
      "source": [
        "### Training loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7YYm33df0w6N",
        "outputId": "72888316-0d2d-431a-96c1-3200127c18ec"
      },
      "source": [
        "metrics = []\n",
        "for perm in perms:\n",
        "    batch = {k: v[perm] for k,v in train_ds.items()}                            #batch is a dictionary/pytree here\n",
        "    optimizer, metric = train_step(optimizer, batch)\n",
        "    metrics.append(metric)\n",
        "\n",
        "metrics = jax.device_get(metrics)                                               #Get metrics from device into CPU as numpy arrays\n",
        "mean_metrics = {k : onp.mean([metric[k] for metric in metrics])                 #Averaging metrics of all batches, while\n",
        "                    for k in metrics[0]}                                        #Looping over all types of metrics\n",
        "print(mean_metrics)                                                             #Can print outside any jit-ted functions"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "{'accuracy': 0.9871461, 'loss': 0.04197854}\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}