ResNet20v1_CIFAR10_Different_Inits.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_Different_Inits.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": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "af9b4dd2-d9d0-41c7-f44b-a7e41c378970"
      },
      "source": [
        "# TensorFlow Imports\n",
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.2.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Z8L9GgmRFd_",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 357
        },
        "outputId": "25856193-5620-405c-ff7e-4e4e4dae525c"
      },
      "source": [
        "# Which GPU?\n",
        "!nvidia-smi"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Wed Jul 15 13:28:13 2020       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 450.51.05    Driver Version: 418.67       CUDA Version: 10.1     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   34C    P8     9W /  70W |      0MiB / 15079MiB |      0%      Default |\n",
            "|                               |                      |                 ERR! |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A-14c7ovy2FX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "!pip install wandb"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VC1vXOBTzCQR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "426f363b-946f-42df-dd43-c3b02dfacdc3"
      },
      "source": [
        "import wandb\n",
        "from wandb.keras import WandbCallback\n",
        "\n",
        "wandb.login()"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/javascript": [
              "\n",
              "        window._wandbApiKey = new Promise((resolve, reject) => {\n",
              "            function loadScript(url) {\n",
              "            return new Promise(function(resolve, reject) {\n",
              "                let newScript = document.createElement(\"script\");\n",
              "                newScript.onerror = reject;\n",
              "                newScript.onload = resolve;\n",
              "                document.body.appendChild(newScript);\n",
              "                newScript.src = url;\n",
              "            });\n",
              "            }\n",
              "            loadScript(\"https://cdn.jsdelivr.net/npm/postmate/build/postmate.min.js\").then(() => {\n",
              "            const iframe = document.createElement('iframe')\n",
              "            iframe.style.cssText = \"width:0;height:0;border:none\"\n",
              "            document.body.appendChild(iframe)\n",
              "            const handshake = new Postmate({\n",
              "                container: iframe,\n",
              "                url: 'https://app.wandb.ai/authorize'\n",
              "            });\n",
              "            const timeout = setTimeout(() => reject(\"Couldn't auto authenticate\"), 5000)\n",
              "            handshake.then(function(child) {\n",
              "                child.on('authorize', data => {\n",
              "                    clearTimeout(timeout)\n",
              "                    resolve(data)\n",
              "                });\n",
              "            });\n",
              "            })\n",
              "        });\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "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": 5,
      "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": {
          "base_uri": "https://localhost:8080/",
          "height": 204
        },
        "outputId": "71d28d11-d9c0-4ada-c2de-08ab135989d0"
      },
      "source": [
        "!wget https://raw.githubusercontent.com/GoogleCloudPlatform/keras-idiomatic-programmer/master/zoo/resnet/resnet_cifar10.py"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-07-15 13:28:32--  https://raw.githubusercontent.com/GoogleCloudPlatform/keras-idiomatic-programmer/master/zoo/resnet/resnet_cifar10.py\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 151.101.0.133, 151.101.64.133, 151.101.128.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|151.101.0.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 6064 (5.9K) [text/plain]\n",
            "Saving to: ‘resnet_cifar10.py’\n",
            "\n",
            "\rresnet_cifar10.py     0%[                    ]       0  --.-KB/s               \rresnet_cifar10.py   100%[===================>]   5.92K  --.-KB/s    in 0s      \n",
            "\n",
            "2020-07-15 13:28:32 (89.5 MB/s) - ‘resnet_cifar10.py’ saved [6064/6064]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "brvfbS3DFr4p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import resnet_cifar10"
      ],
      "execution_count": 7,
      "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": 8,
      "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": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "98dff665-4d64-4b04-da0c-4b1ee3abd67c"
      },
      "source": [
        "# Load the training set of CIFAR10\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading data from https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz\n",
            "170500096/170498071 [==============================] - 6s 0us/step\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "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",
        "# 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": 11,
      "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": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-EMm1uWXdhPu",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "outputId": "3d94719c-3252-4353-ed26-1c3dd96edf3d"
      },
      "source": [
        "rng = rng = [i for i in range(40)]\n",
        "plt.plot([lr_schedule(x) for x in rng])\n",
        "plt.show()"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYYAAAD4CAYAAADo30HgAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAdpklEQVR4nO3dfZBd9X3f8fdnH7Ur0C5arTGWtFphrewsNn7IWnYcQ+JgB+GkKA/QSHUzykSuOh5o09DUFXVLiTLqmKQxbcYQj1IojOtYqMSOdxLZ+AGPsacgIWywJbBgw6Nkg2QhCWy0klb77R/3t9K9V3e1Zx+053Lv5zWj0bm/8zvnfs8B7Wfv+Z37O4oIzMzMxjTkXYCZmVUXB4OZmZVwMJiZWQkHg5mZlXAwmJlZiaa8C5gJCxYsiN7e3rzLMDN7XXnkkUd+GhHd5e01EQy9vb3s3Lkz7zLMzF5XJD1Xqd2XkszMrISDwczMSjgYzMyshIPBzMxKOBjMzKxEpmCQtFLSHklDkjZUWN8q6Z60fruk3qJ1N6b2PZKuLGq/U9J+Sbsq7O/fSPqRpN2S/nxqh2ZmZlMxYTBIagRuA64C+oE1kvrLuq0DDkXEMuBW4Ja0bT+wGrgEWAncnvYHcFdqK3+/DwKrgHdExCXAf5/8YZmZ2VRl+R7DCmAoIp4GkLSFwg/ux4v6rAJuTsv3Ap+RpNS+JSKOAc9IGkr7ezAiHij+ZFHk48Cn0jZExP7JHlRW33ziJR574fC52n1Va2lq4Pd/qZeOtua8SzGzKpMlGBYCLxS93gu8d7w+ETEi6QjQldofKtt24QTvtxy4TNImYBj4k4h4uLyTpPXAeoCenp4Mh3Gmbz95gM89VPH7HTVt7BEcF3W08bu/uCjfYsys6lTjN5+bgPnA+4D3AFslXRxlTxSKiM3AZoCBgYEpPW1o46q3sXHV26ZZ7uvPkddO8I6NX+PQa8fzLsXMqlCWwed9wOKi14tSW8U+kpqADuBgxm3L7QW+GAU7gFFgQYY6LaPz5zQhwStHT+RdiplVoSzB8DDQJ2mppBYKg8mDZX0GgbVp+Rrg/vQb/iCwOt21tBToA3ZM8H5/D3wQQNJyoAX4aZaDsWwaGkRHWzOHHQxmVsGEwRARI8D1wH3AE8DWiNgtaaOkq1O3O4CuNLh8A7Ahbbsb2EphoPqrwHURcRJA0heAB4G3SNoraV3a153Axek21i3A2vLLSDZ9HW3NHH7NwWBmZ8o0xhAR24BtZW03FS0PA9eOs+0mYFOF9jXj9D8O/MssddnUdfoTg5mNw998rlMd7S0ccTCYWQUOhjrV0dbMEd+VZGYVOBjqlC8lmdl4HAx1qrO9mVeOnmB01OP6ZlbKwVCnOtqaGQ149dhI3qWYWZVxMNSpsTmSjviWVTMr42CoU53tLQC+M8nMzuBgqFOd7YVPDIeP+s4kMyvlYKhTY5eS/O1nMyvnYKhTnWNjDL6UZGZlHAx1ap6DwczG4WCoU3OaG5nT3MBhf/vZzMo4GOpYZ5vnSzKzMzkY6lhnu6feNrMzORjq2DzPl2RmFTgY6lhnW7Mf72lmZ3Aw1DFfSjKzShwMdazw3GfflWRmpRwMdayzvYXhE6MMnziZdylmVkUyBYOklZL2SBqStKHC+lZJ96T12yX1Fq27MbXvkXRlUfudkvZL2jXOe/57SSFpweQPy7IYmxbD4wxmVmzCYJDUCNwGXAX0A2sk9Zd1WwcciohlwK3ALWnbfmA1cAmwErg97Q/grtRW6T0XA78OPD/J47FJODVfkoPBzIpk+cSwAhiKiKcj4jiwBVhV1mcVcHdavhe4QpJS+5aIOBYRzwBDaX9ExAPAy+O8563AJwA/XuwcOjXDqgegzaxIlmBYCLxQ9HpvaqvYJyJGgCNAV8ZtS0haBeyLiMcm6Lde0k5JOw8cOJDhMKxcZ5ufyWBmZ6qqwWdJ7cB/Am6aqG9EbI6IgYgY6O7uPvfF1aDTU2/7ziQzOy1LMOwDFhe9XpTaKvaR1AR0AAczblvszcBS4DFJz6b+35P0xgx12iR1tHuGVTM7U5ZgeBjok7RUUguFweTBsj6DwNq0fA1wf0REal+d7lpaCvQBO8Z7o4j4YUS8ISJ6I6KXwqWnd0fEi5M6Ksvk/NYmGuRgMLNSEwZDGjO4HrgPeALYGhG7JW2UdHXqdgfQJWkIuAHYkLbdDWwFHge+ClwXEScBJH0BeBB4i6S9ktbN7KHZRBoaVJgvyYPPZlakKUuniNgGbCtru6loeRi4dpxtNwGbKrSvyfC+vVnqs6nr9ER6ZlamqgafbfZ1tPuZDGZWysFQ5zrbmjniu5LMrIiDoc51+FKSmZVxMNS5zvZmX0oysxIOhjrX2VYIhtFRzz5iZgUOhjo3r62ZCHh1eCTvUsysSjgY6lxnu+dLMrNSDoY613lq6m3fmWRmBQ6GOtfhqbfNrIyDoc6NfWLwpSQzG+NgqHOnPjE4GMwscTDUubFnMvjbz2Y2xsFQ51qbGmlrbvQYg5md4mAwf/vZzEo4GMzzJZlZCQeD0dHWzBFfSjKzxMFgvpRkZiUcDJYuJfmuJDMrcDAYne0tvivJzE5xMBgdbc0cGxll+MTJvEsxsyqQKRgkrZS0R9KQpA0V1rdKuiet3y6pt2jdjal9j6Qri9rvlLRf0q6yff2FpB9J+oGkL0nqnPrhWRYdnhbDzIpMGAySGoHbgKuAfmCNpP6ybuuAQxGxDLgVuCVt2w+sBi4BVgK3p/0B3JXayn0deFtEXAo8Cdw4yWOySer0RHpmViTLJ4YVwFBEPB0Rx4EtwKqyPquAu9PyvcAVkpTat0TEsYh4BhhK+yMiHgBeLn+ziPhaRIw9NeYhYNEkj8kmqbPNz2Qws9OyBMNC4IWi13tTW8U+6Yf6EaAr47Zn84fAVyqtkLRe0k5JOw8cODCJXVq5058YfGeSmVXx4LOkTwIjwOcrrY+IzRExEBED3d3ds1tcjelo8wyrZnZalmDYBywuer0otVXsI6kJ6AAOZtz2DJL+APhN4KMR4afUn2NjU2+/4mAwM7IFw8NAn6SlklooDCYPlvUZBNam5WuA+9MP9EFgdbpraSnQB+w425tJWgl8Arg6Il7Lfig2Vee3NtHYIA8+mxmQIRjSmMH1wH3AE8DWiNgtaaOkq1O3O4AuSUPADcCGtO1uYCvwOPBV4LqIOAkg6QvAg8BbJO2VtC7t6zPA+cDXJT0q6bMzdKw2DknMm9Pkbz+bGQBNWTpFxDZgW1nbTUXLw8C142y7CdhUoX3NOP2XZanJZlZnewtHjo5M3NHMal7VDj7b7Opoa/ZdSWYGOBgs6WjzDKtmVuBgMMBTb5vZaQ4GA6Czrdl3JZkZ4GCwpKOtmVeGT3By1F8bMat3DgYDoKO9hQh4ddifGszqnYPBgMKlJPBEembmYLDk1HxJHmcwq3sOBgOKZlj1JwazuudgMOB0MPhSkpk5GAyAeWNjDP72s1ndczAY4DEGMzvNwWAAtDY10t7S6EtJZuZgsNM625o9+GxmDgY7bZ6nxTAzHAxWpLO92Y/3NDMHg53W2dbip7iZmYPBTuvwpSQzw8FgRfxMBjODjMEgaaWkPZKGJG2osL5V0j1p/XZJvUXrbkzteyRdWdR+p6T9knaV7Wu+pK9Leir9fcHUD88mo6O9mWMjowyfOJl3KWaWowmDQVIjcBtwFdAPrJHUX9ZtHXAoIpYBtwK3pG37gdXAJcBK4Pa0P4C7Ulu5DcA3I6IP+GZ6bbPAX3IzM8j2iWEFMBQRT0fEcWALsKqszyrg7rR8L3CFJKX2LRFxLCKeAYbS/oiIB4CXK7xf8b7uBn5rEsdj09DZ1gJ4viSzepclGBYCLxS93pvaKvaJiBHgCNCVcdtyF0bET9Lyi8CFlTpJWi9pp6SdBw4cyHAYNpFTM6x6viSzulbVg88REUDFZ01GxOaIGIiIge7u7lmurDadupTkTwxmdS1LMOwDFhe9XpTaKvaR1AR0AAczblvuJUkXpX1dBOzPUKPNgA4/xc3MyBYMDwN9kpZKaqEwmDxY1mcQWJuWrwHuT7/tDwKr011LS4E+YMcE71e8r7XAlzPUaDPg1DMZPPhsVtcmDIY0ZnA9cB/wBLA1InZL2ijp6tTtDqBL0hBwA+lOoojYDWwFHge+ClwXEScBJH0BeBB4i6S9ktalfX0K+LCkp4APpdc2C85rbaKxQf72s1mda8rSKSK2AdvK2m4qWh4Grh1n203Apgrta8bpfxC4IktdNrMk+dvPZlbdg882+zrb/O1ns3rnYLAS8xwMZnXPwWAlOtt9Kcms3jkYrIQvJZmZg8FKFAaffVeSWT1zMFiJjvYWXhke4eRoxS+cm1kdcDBYic707edXh305yaxeORisxOmJ9BwMZvXKwWAlPJGemTkYrMSp+ZIcDGZ1y8FgJTrSw3p8Z5JZ/XIwWAlPvW1mDgYrcSoYPPhsVrccDFaipamBuS2NHnw2q2MOBjuDp942q28OBjtDR3uLxxjM6piDwc5QmEjPdyWZ1SsHg53Bl5LM6puDwc7Q2d7swWezOpYpGCStlLRH0pCkDRXWt0q6J63fLqm3aN2NqX2PpCsn2qekKyR9T9Kjkr4radn0DtEmq6O98EyGCM+walaPJgwGSY3AbcBVQD+wRlJ/Wbd1wKGIWAbcCtyStu0HVgOXACuB2yU1TrDPvwY+GhHvBP4W+M/TO0SbrI62Zo6PjDJ8YjTvUswsB00Z+qwAhiLiaQBJW4BVwONFfVYBN6fle4HPSFJq3xIRx4BnJA2l/XGWfQYwL/XpAH48tUOzqepM02L82T8+TntzY87VVJemxgbWfWAp3ee35l2K2TmTJRgWAi8Uvd4LvHe8PhExIukI0JXaHyrbdmFaHm+fHwO2SToKvAK8r1JRktYD6wF6enoyHIZl9baF85g/t4Uvf39f3qVUlQBeO36SBee18LHLLs67HLNzJkswzLY/Bj4SEdsl/Qfg0xTCokREbAY2AwwMDPhi+Ay6dFEn3/svH867jKoTEVz6p1/j+Zdfy7sUs3Mqy+DzPmBx0etFqa1iH0lNFC4BHTzLthXbJXUD74iI7an9HuD9mY7E7ByTxJKudp496GCw2pYlGB4G+iQtldRCYTB5sKzPILA2LV8D3B+FW1oGgdXprqWlQB+w4yz7PAR0SFqe9vVh4ImpH57ZzFoyfy7PH/x53mWYnVMTXkpKYwbXA/cBjcCdEbFb0kZgZ0QMAncAn0uDyy9T+EFP6reVwqDyCHBdRJwEqLTP1P6vgL+TNEohKP5wRo/YbBqWdLVz3+4XGTk5SlOjvwZktSnTGENEbAO2lbXdVLQ8DFw7zrabgE1Z9pnavwR8KUtdZrNtSVc7I6PBjw8P09PVnnc5ZueEf+Uxm4QlXXMBeO5lX06y2uVgMJuEJelTggegrZY5GMwm4cLz59DS1OABaKtpDgazSWhoEEvmt/OcPzFYDXMwmE3Ski4Hg9U2B4PZJPXMn8tzL//cs89azXIwmE1S74J2hk+Msv/VY3mXYnZOOBjMJqlnfuHOJF9OslrlYDCbpN6x7zL4ziSrUQ4Gs0laeEEbjQ3yJwarWQ4Gs0lqbmzgTZ1zeM7Tb1uNcjCYTUFvl2dZtdrlYDCbgp75fi6D1S4Hg9kULOlq58jRExx+7XjepZjNOAeD2RScmmXVnxqsBjkYzKZgbJZVD0BbLXIwmE3B2JfcPABttcjBYDYF7S1NvOH8Vg9AW01yMJhN0ZKudp53MFgNyhQMklZK2iNpSNKGCutbJd2T1m+X1Fu07sbUvkfSlRPtUwWbJD0p6QlJ/3Z6h2h2bizpmutHfFpNmjAYJDUCtwFXAf3AGkn9Zd3WAYciYhlwK3BL2rYfWA1cAqwEbpfUOME+/wBYDLw1In4B2DKtIzQ7R5bMb+elV45x9PjJvEsxm1FZPjGsAIYi4umIOE7hB/Wqsj6rgLvT8r3AFZKU2rdExLGIeAYYSvs72z4/DmyMiFGAiNg/9cMzO3d60p1Jz/vOJKsxWYJhIfBC0eu9qa1in4gYAY4AXWfZ9mz7fDPwe5J2SvqKpL5KRUlan/rsPHDgQIbDMJtZY7OsPus7k6zGVOPgcyswHBEDwN8Ad1bqFBGbI2IgIga6u7tntUAzOP1dBg9AW63JEgz7KFzzH7MotVXsI6kJ6AAOnmXbs+1zL/DFtPwl4NIMNZrNus72Fjramj0AbTUnSzA8DPRJWiqphcJg8mBZn0FgbVq+Brg/Cg/EHQRWp7uWlgJ9wI4J9vn3wAfT8q8AT07t0MzOvSVd7Z4Ww2pO00QdImJE0vXAfUAjcGdE7Ja0EdgZEYPAHcDnJA0BL1P4QU/qtxV4HBgBrouIkwCV9pne8lPA5yX9MfAz4GMzd7hmM6tnfjs/2Hsk7zLMZtSEwQAQEduAbWVtNxUtDwPXjrPtJmBTln2m9sPAb2SpyyxvvV1z+cquFzlxcpTmxmocsjObPP+fbDYNPV3tnBwN9h06mncpZjPGwWA2DUvme5ZVqz0OBrNp6F0w9lwG35lktcPBYDYNbzi/lTnNDb4zyWqKg8FsGiSxZP5cB4PVFAeD2TT1dLX7UpLVFAeD2TQtmd/O8y+/xuho5F2K2YxwMJhN05IFczk2Msr+V4/lXYrZjHAwmE3T2C2rnmXVaoWDwWyaPMuq1RoHg9k0Lexso6lB/sRgNcPBYDZNTY0NLLygzd9+tprhYDCbAT3z230pyWqGg8FsBvR2zeXZgz+n8BgSs9c3B4PZDFjS1c6rwyMcfu1E3qWYTZuDwWwGLOlKk+l5nMFqgIPBbAaM3bLqqTGsFjgYzGZAz9hzGTwAbTXAwWA2A+Y0N/LGeXP8XQarCQ4GsxnS0+VbVq02ZAoGSSsl7ZE0JGlDhfWtku5J67dL6i1ad2Nq3yPpykns868k/Wxqh2U2+5bMb/fgs9WEpok6SGoEbgM+DOwFHpY0GBGPF3VbBxyKiGWSVgO3AL8nqR9YDVwCvAn4hqTlaZtx9ylpALhgRo7QbJb0LpjLgUf28uPDR2lrbsy7nKrS0CA62przLsMymjAYgBXAUEQ8DSBpC7AKKA6GVcDNafle4DOSlNq3RMQx4BlJQ2l/jLfPFER/AfwL4LencWxms2ppev7z+z91f86VVKe/vPYd/O4vLsq7DMsgSzAsBF4oer0XeO94fSJiRNIRoCu1P1S27cK0PN4+rwcGI+InhWypTNJ6YD1AT09PhsMwO7eu+IU38KnfeTvDJ07mXUrV+ey3n+Yru150MLxOZAmGWSPpTcC1wK9O1DciNgObAQYGBjwPgeWutamR1Sv8S0olT+7/GV/+/j6Oj4zS0uR7Xqpdlv9C+4DFRa8XpbaKfSQ1AR3AwbNsO177u4BlwJCkZ4H2dPnJzF7HLu/r5ufHT/K95w/lXYplkCUYHgb6JC2V1EJhMHmwrM8gsDYtXwPcH4XZxAaB1emupaVAH7BjvH1GxD9GxBsjojcieoHXImLZdA/SzPL1/mVdNDaIB548kHcplsGEwRARIxSu+98HPAFsjYjdkjZKujp1uwPoSr/d3wBsSNvuBrZSGKj+KnBdRJwcb58ze2hmVi3mzWnmXYs7+c5TP827FMtAtTBN8MDAQOzcuTPvMszsLP7qm09x6zeeZOcnP0TXea15l2OApEciYqC83aNAZjYrLl/eTQR8d8ifGqqdg8HMZsXbF3bQ2d7MA086GKqdg8HMZkVjg/jlZQv4zlMH/KS7KudgMLNZ8yt93ex/9Rg/evHVvEuxs3AwmNmsuWz5AgDftlrlHAxmNmsu6mhj+YXn+bbVKudgMLNZdVlfNzuefZmjxz2nVLVyMJjZrLp8eTfHR0Z56JmDeZdi43AwmNmseu/S+bQ2NXicoYo5GMxsVs1pbmTF0vkeZ6hiDgYzm3WX93UztP9n/Pjw0bxLsQocDGY26y5f3g34ttVq5WAws1m3/MLzeOO8OTzwlIOhGjkYzGzWSeKyvgV896mfcnLU02NUGweDmeXisuXdvDI8wmN7D+ddipVxMJhZLi5btgDJ4wzVyMFgZrm4YG4Lly7scDBUIQeDmeXmsr5uHn3hMEeOnsi7FCviYDCz3Fy+vJvRgP/np7pVlUzBIGmlpD2ShiRtqLC+VdI9af12Sb1F625M7XskXTnRPiV9PrXvknSnpObpHaKZVat39XRyXmuTb1utMhMGg6RG4DbgKqAfWCOpv6zbOuBQRCwDbgVuSdv2A6uBS4CVwO2SGifY5+eBtwJvB9qAj03rCM2sajU3NvD+N3fxwJM/9VPdqkhThj4rgKGIeBpA0hZgFfB4UZ9VwM1p+V7gM5KU2rdExDHgGUlDaX+Mt8+I2Da2U0k7gEVTPDYzex24bHk3X3v8JT706W/TIOVdzuvOf/udt/Oe3vkzus8swbAQeKHo9V7gveP1iYgRSUeArtT+UNm2C9PyWfeZLiH9PvBHlYqStB5YD9DT05PhMMysGv2zSy/i+88fYviEn88wFW3NjTO+zyzBkJfbgQci4juVVkbEZmAzwMDAgD+Dmr1Odba38Ol//s68y7AiWYJhH7C46PWi1Fapz15JTUAHcHCCbcfdp6T/CnQD/zpDfWZmNoOy3JX0MNAnaamkFgqDyYNlfQaBtWn5GuD+KIwkDQKr011LS4E+YMfZ9inpY8CVwJqIGJ3e4ZmZ2WRN+IkhjRlcD9wHNAJ3RsRuSRuBnRExCNwBfC4NLr9M4Qc9qd9WCgPVI8B1EXESoNI+01t+FngOeLAwfs0XI2LjjB2xmZmdlWrhFrGBgYHYuXNn3mWYmb2uSHokIgbK2/3NZzMzK+FgMDOzEg4GMzMr4WAwM7MSNTH4LOkAhTuZpmIBUK1TO7q2qXFtU+Papub1XNuSiOgub6yJYJgOSTsrjcpXA9c2Na5talzb1NRibb6UZGZmJRwMZmZWwsGQJuKrUq5talzb1Li2qam52up+jMHMzEr5E4OZmZVwMJiZWYm6DgZJKyXtkTQkaUPe9RST9KykH0p6VFKuMwRKulPSfkm7itrmS/q6pKfS3xdUUW03S9qXzt2jkj6SU22LJX1L0uOSdkv6o9Se+7k7S225nztJcyTtkPRYqu1PU/tSSdvTv9d70pT91VLbXZKeKTpvuT15SFKjpO9L+of0evLnLSLq8g+F6b7/CbgYaAEeA/rzrquovmeBBXnXkWq5HHg3sKuo7c+BDWl5A3BLFdV2M/AnVXDeLgLenZbPB54E+qvh3J2lttzPHSDgvLTcDGwH3gdsBVan9s8CH6+i2u4Crsn7/7lU1w3A3wL/kF5P+rzV8yeGFcBQRDwdEceBLcCqnGuqShHxAIXnbBRbBdydlu8GfmtWi0rGqa0qRMRPIuJ7aflV4AkKzzzP/dydpbbcRcHP0svm9CeAXwPuTe15nbfxaqsKkhYBvwH8r/RaTOG81XMwLAReKHq9lyr5h5EE8DVJj0han3cxFVwYET9Jyy8CF+ZZTAXXS/pButSUy2WuYpJ6gXdR+A2zqs5dWW1QBecuXQ55FNgPfJ3Cp/vDETGSuuT277W8togYO2+b0nm7VVJrHrUB/wP4BDD29MsupnDe6jkYqt0HIuLdwFXAdZIuz7ug8UThM2rV/NYE/DXwZuCdwE+Av8yzGEnnAX8H/LuIeKV4Xd7nrkJtVXHuIuJkRLyTwvPgVwBvzaOOSsprk/Q24EYKNb4HmA/8x9muS9JvAvsj4pHp7queg2EfsLjo9aLUVhUiYl/6ez/wJQr/OKrJS5IuAkh/78+5nlMi4qX0j3cU+BtyPHeSmin84P18RHwxNVfFuatUWzWdu1TPYeBbwC8BnZLGHkec+7/XotpWpktzERHHgP9NPuftl4GrJT1L4dL4rwH/kymct3oOhoeBvjRi30LhOdWDOdcEgKS5ks4fWwZ+Hdh19q1m3SCwNi2vBb6cYy0lxn7oJr9NTucuXd+9A3giIj5dtCr3czdebdVw7iR1S+pMy23AhymMgXwLuCZ1y+u8VartR0VBLwrX8Gf9vEXEjRGxKCJ6Kfw8uz8iPspUzlveI+h5/gE+QuFujH8CPpl3PUV1XUzhLqnHgN151wZ8gcJlhRMUrlGuo3Dt8pvAU8A3gPlVVNvngB8CP6DwQ/iinGr7AIXLRD8AHk1/PlIN5+4steV+7oBLge+nGnYBN6X2i4EdwBDwf4HWKqrt/nTedgH/h3TnUl5/gF/l9F1Jkz5vnhLDzMxK1POlJDMzq8DBYGZmJRwMZmZWwsFgZmYlHAxmZlbCwWBmZiUcDGZmVuL/A6M+qPk/BTt9AAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "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": 14,
      "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": 15,
      "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": 16,
      "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": [
        "for i in range(10):\n",
        "    id = 'resnet20v1-aug-diff-inits-corrected' + str(i)\n",
        "    wandb.init(project='loss-landscape', id=id)\n",
        "    \n",
        "    # Train model\n",
        "    model = get_training_model()\n",
        "    model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
        "    \n",
        "    start = time.time()\n",
        "    model.fit(train_ds,\n",
        "            validation_data=test_ds,\n",
        "            epochs=40,\n",
        "            callbacks=[lr_callback, WandbCallback(), cm_callback])\n",
        "    end = time.time()\n",
        "    print(\"Network takes {:.3f} seconds to train\".format(end - start))\n",
        "    \n",
        "    wandb.log({'training_time': end - start})\n",
        "    wandb.log({'nb_model_params': model.count_params()})\n",
        "    model_name = 'resnet20v1-aug-diff-inits-' + str(i) + '.h5'\n",
        "    model.save(model_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZlXvPC5VRvU3",
        "colab_type": "text"
      },
      "source": [
        "Outputs have been Omitted for brevity. "
      ]
    },
    {
      "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": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NIrhJ6um1J6e",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 221
        },
        "outputId": "9b24e6fe-a931-4a17-dab3-032578f95177"
      },
      "source": [
        "!gsutil -m cp -r resnet20v1-aug-diff-inits-*.h5 gs://losslandscape/ResNet20v1_CIFAR10_Aug_Diff_Inits/"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying file://resnet20v1-aug-diff-inits-0.h5 [Content-Type=application/octet-stream]...\n",
            "/ [0/10 files][    0.0 B/ 69.4 MiB]   0% Done                                   \rCopying file://resnet20v1-aug-diff-inits-7.h5 [Content-Type=application/octet-stream]...\n",
            "/ [0/10 files][    0.0 B/ 69.4 MiB]   0% Done                                   \rCopying file://resnet20v1-aug-diff-inits-1.h5 [Content-Type=application/octet-stream]...\n",
            "/ [0/10 files][    0.0 B/ 69.4 MiB]   0% Done                                   \rCopying file://resnet20v1-aug-diff-inits-6.h5 [Content-Type=application/octet-stream]...\n",
            "/ [0/10 files][    0.0 B/ 69.4 MiB]   0% Done                                   \rCopying file://resnet20v1-aug-diff-inits-5.h5 [Content-Type=application/octet-stream]...\n",
            "/ [0/10 files][    0.0 B/ 69.4 MiB]   0% Done                                   \rCopying file://resnet20v1-aug-diff-inits-4.h5 [Content-Type=application/octet-stream]...\n",
            "Copying file://resnet20v1-aug-diff-inits-9.h5 [Content-Type=application/octet-stream]...\n",
            "Copying file://resnet20v1-aug-diff-inits-8.h5 [Content-Type=application/octet-stream]...\n",
            "Copying file://resnet20v1-aug-diff-inits-2.h5 [Content-Type=application/octet-stream]...\n",
            "Copying file://resnet20v1-aug-diff-inits-3.h5 [Content-Type=application/octet-stream]...\n",
            "-\n",
            "Operation completed over 10 objects/69.4 MiB.                                    \n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}