ResNet20v1_CIFAR10.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "scratchpad",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ayulockin/LossLandscape/blob/master/ResNet20v1_CIFAR10.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YDmOU9DezWrm",
        "colab_type": "text"
      },
      "source": [
        "## Set up and imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lIYdn1woOS1n",
        "colab": {}
      },
      "source": [
        "# TensorFlow Imports\n",
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Z8L9GgmRFd_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Which GPU?\n",
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A-14c7ovy2FX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "!pip install wandb"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VC1vXOBTzCQR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import wandb\n",
        "from wandb.keras import WandbCallback\n",
        "\n",
        "wandb.login()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "txX7OoEpROgI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Other imports\n",
        "from tensorflow.keras.layers import *\n",
        "from tensorflow.keras.models import *\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import time\n",
        "\n",
        "# Random seed fixation\n",
        "tf.random.set_seed(666)\n",
        "np.random.seed(666)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GlDedsctzZBE",
        "colab_type": "text"
      },
      "source": [
        "## Get the model from [keras-idiomatic-programmer](https://github.com/GoogleCloudPlatform/keras-idiomatic-programmer)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5-_4ROoYRHJq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget https://raw.githubusercontent.com/GoogleCloudPlatform/keras-idiomatic-programmer/master/zoo/resnet/resnet_cifar10.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "brvfbS3DFr4p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import resnet_cifar10"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8kEQH0sAzn9-",
        "colab_type": "text"
      },
      "source": [
        "## Utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1NlIHQ-rRWlX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_training_model():\n",
        "    # ResNet20\n",
        "    n = 2\n",
        "    depth =  n * 9 + 2\n",
        "    n_blocks = ((depth - 2) // 9) - 1\n",
        "\n",
        "    # The input tensor\n",
        "    inputs = Input(shape=(32, 32, 3))\n",
        "\n",
        "    # The Stem Convolution Group\n",
        "    x = resnet_cifar10.stem(inputs)\n",
        "\n",
        "    # The learner\n",
        "    x = resnet_cifar10.learner(x, n_blocks)\n",
        "\n",
        "    # The Classifier for 10 classes\n",
        "    outputs = resnet_cifar10.classifier(x, 10)\n",
        "\n",
        "    # Instantiate the Model\n",
        "    model = Model(inputs, outputs)\n",
        "    \n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7h-p6STrz14i",
        "colab_type": "text"
      },
      "source": [
        "## Construct data loaders"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jp-VUHKiRqo4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the training set of CIFAR10\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dNhaSmKjR2Vq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "BATCH_SIZE = 128\n",
        "\n",
        "def normalize(image, label):\n",
        "    return tf.image.convert_image_dtype(image, tf.float32), label\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M01oYoA7z5jI",
        "colab_type": "text"
      },
      "source": [
        "## Model sanity checks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ymnw4EirV8HT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = get_training_model()\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IF-7vpPdz8Bb",
        "colab_type": "text"
      },
      "source": [
        "## Callbacks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cZWsbIFPV4-2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Custom LR schedule as mentioned in the LossLandscape paper\n",
        "LR_SCHEDULE = [\n",
        "    # (epoch to start, learning rate) tuples\n",
        "    (0, 1.6*1e-3),\n",
        "    (9, (1.6*1e-3)/2),\n",
        "    (19, (1.6*1e-3)/4),\n",
        "    (29, (1.6*1e-3)/8),\n",
        "]\n",
        "\n",
        "def lr_schedule(epoch):\n",
        "    if (epoch >= 0) & (epoch < 9):\n",
        "        return LR_SCHEDULE[0][1]\n",
        "    elif (epoch >= 9) & (epoch < 19):\n",
        "        return LR_SCHEDULE[1][1]\n",
        "    elif (epoch >= 19) & (epoch < 29):\n",
        "        return LR_SCHEDULE[2][1]\n",
        "    else:\n",
        "        return LR_SCHEDULE[3][1]\n",
        "\n",
        "lr_callback = tf.keras.callbacks.LearningRateScheduler(lambda epoch: lr_schedule(epoch), verbose=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-EMm1uWXdhPu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "rng = rng = [i for i in range(40)]\n",
        "plt.plot([lr_schedule(x) for x in rng])\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fuZGnXYbyyAl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "SAVE_PATH = '/content/ResNet20v1_CIFAR10'\n",
        "\n",
        "def save_model(epoch, logs):\n",
        "    model.save(SAVE_PATH+'resnet20v1_checkpoint_{}.h5'.format(epoch))\n",
        "\n",
        "save_callback = tf.keras.callbacks.LambdaCallback(on_epoch_end=save_model, verbose=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X-5U9pxv5ydj",
        "colab_type": "text"
      },
      "source": [
        "A custom callback to log confusion matrix batchwise (referred from this [tutorial](https://www.tensorflow.org/tensorboard/image_summaries)). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KrYoLvtT3_l9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from sklearn.metrics import confusion_matrix\n",
        "import itertools\n",
        "import io"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZJdGjP3Z5ts4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "CLASS_NAMES = [\"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\", \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tAcwuGvr3hKy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def plot_confusion_matrix(cm, class_names):\n",
        "    figure = plt.figure(figsize=(8, 8))\n",
        "    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)\n",
        "    plt.title(\"Confusion matrix\")\n",
        "    plt.colorbar()\n",
        "    tick_marks = np.arange(len(class_names))\n",
        "    plt.xticks(tick_marks, class_names, rotation=45)\n",
        "    plt.yticks(tick_marks, class_names)\n",
        "\n",
        "    # Normalize the confusion matrix.\n",
        "    cm = np.around(cm.astype('float') / cm.sum(axis=1)[:, np.newaxis], decimals=2)\n",
        "\n",
        "    # Use white text if squares are dark; otherwise black.\n",
        "    threshold = cm.max() / 2.\n",
        "    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):\n",
        "        color = \"white\" if cm[i, j] > threshold else \"black\"\n",
        "        plt.text(j, i, cm[i, j], horizontalalignment=\"center\", color=color)\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.ylabel('True label')\n",
        "    plt.xlabel('Predicted label')\n",
        "    \n",
        "    return figure\n",
        "\n",
        "def plot_to_image(figure):\n",
        "    \"\"\"Converts the matplotlib plot specified by 'figure' to a PNG image and\n",
        "    returns it. The supplied figure is closed and inaccessible after this call.\"\"\"\n",
        "    # Save the plot to a PNG in memory.\n",
        "    buf = io.BytesIO()\n",
        "    plt.savefig(buf, format='png')\n",
        "    # Closing the figure prevents it from being displayed directly inside\n",
        "    # the notebook.\n",
        "    plt.close(figure)\n",
        "    buf.seek(0)\n",
        "    # Convert PNG buffer to TF image\n",
        "    image = tf.image.decode_png(buf.getvalue(), channels=3)\n",
        "    # Convert back to NumPy\n",
        "    image = image.numpy()\n",
        "    return image\n",
        "\n",
        "def log_confusion_matrix(epoch, logs):\n",
        "    # Use the model to predict the values from the validation dataset\n",
        "    test_pred_raw = model.predict(x_test)\n",
        "    test_pred = np.argmax(test_pred_raw, axis=1)\n",
        "\n",
        "    # Calculate the confusion matrix\n",
        "    cm = confusion_matrix(y_test, test_pred)\n",
        "    # Log the confusion matrix as an image to wandb\n",
        "    figure = plot_confusion_matrix(cm, class_names=CLASS_NAMES)\n",
        "    cm_image = plot_to_image(figure)\n",
        "    wandb.log({'confusion_matrix': wandb.Image(cm_image)})\n",
        "\n",
        "cm_callback = tf.keras.callbacks.LambdaCallback(on_epoch_end=log_confusion_matrix, verbose=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m0XcW9Yv0q-o",
        "colab_type": "text"
      },
      "source": [
        "## Model training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WOv6rFsbR9z2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "wandb.init(project='loss-landscape', id='resnet20v1-no-aug-1')\n",
        "\n",
        "# Train model\n",
        "model = get_training_model()\n",
        "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=40,\n",
        "         callbacks=[lr_callback, WandbCallback(), cm_callback, save_callback])\n",
        "\n",
        "end = time.time()\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))\n",
        "wandb.log({'training_time': end - start})\n",
        "wandb.log({'nb_model_params': model.count_params()})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hGjUF8bEulXV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Serialize the model\n",
        "model.save('resnet20v1_cifar10_40epochs.h5')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Umq0Hyd6xVPQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Adding data augmentation\n",
        "def augment(image,label):\n",
        "    image = tf.image.resize_with_crop_or_pad(image, 40, 40) # Add 8 pixels of padding\n",
        "    image = tf.image.random_crop(image, size=[32, 32, 3]) # Random crop back to 32x32\n",
        "    image = tf.image.random_brightness(image, max_delta=0.5) # Random brightness\n",
        "    image = tf.clip_by_value(image, 0., 1.)\n",
        "    \n",
        "    return image, label\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UcGr4s4DAifS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "SAVE_PATH = '/content/ResNet20v1_CIFAR10_Aug'\n",
        "\n",
        "def save_model(epoch, logs):\n",
        "    model.save(SAVE_PATH+'resnet20v1_checkpoint_{}.h5'.format(epoch))\n",
        "\n",
        "save_callback = tf.keras.callbacks.LambdaCallback(on_epoch_end=save_model, verbose=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5DrzQlXOxgCy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "wandb.init(project='loss-landscape', id='resnet20v1-aug')\n",
        "\n",
        "# Train model\n",
        "model = get_training_model()\n",
        "model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=40,\n",
        "         callbacks=[lr_callback, WandbCallback(), cm_callback, save_callback])\n",
        "\n",
        "end = time.time()\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))\n",
        "wandb.log({'training_time': end - start})\n",
        "wandb.log({'nb_model_params': model.count_params()})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mah_NGL6xilE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Serialize the model\n",
        "model.save('resnet20v1_cifar10_40epochs_data_aug.h5')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8KSGyTxHx1HR",
        "colab_type": "text"
      },
      "source": [
        "## Put the model weights in a GCS bucket"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jWSj_N_dx0D0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import auth as google_auth\n",
        "google_auth.authenticate_user()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NIrhJ6um1J6e",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gsutil -m cp -r ResNet20v1_CIFAR10resnet20v1_checkpoint_*.h5 gs://losslandscape/ResNet20v1_CIFAR10/\n",
        "!gsutil cp resnet20v1_cifar10_40epochs.h5 gs://losslandscape/ResNet20v1_CIFAR10/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LOFfx36KE9kx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gsutil -m cp -r ResNet20v1_CIFAR10_Augresnet20v1_checkpoint_*.h5 gs://losslandscape/ResNet20v1_CIFAR10_Aug/\n",
        "!gsutil cp resnet20v1_cifar10_40epochs_data_aug.h5 gs://losslandscape/ResNet20v1_CIFAR10_Aug/"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}