Calculate_Epsilon.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Calculate_Epsilon.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "environment": {
      "name": "tf2-gpu.2-4.mnightly-2021-01-20-debian-10-test",
      "type": "gcloud",
      "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-4:mnightly-2021-01-20-debian-10-test"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.6"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/Generalized-ODIN-TF/blob/main/Calculate_Epsilon.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QoqcWOsG8uQw"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hSP_ddoJ6FSy"
      },
      "source": [
        "# Grab the initial model weights\n",
        "!wget -q https://github.com/sayakpaul/Generalized-ODIN-TF/releases/download/v1.0.0/models.tar.gz\n",
        "!tar xf models.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3HpaHaAISdK3"
      },
      "source": [
        "!git clone https://github.com/sayakpaul/Generalized-ODIN-TF"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hjNMaqP16fUa"
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "tf.random.set_seed(42)\n",
        "np.random.seed(42)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sa-bFmFe8xz8"
      },
      "source": [
        "## Load the pre-trained model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aSfgIZ1y6qQ6",
        "outputId": "195a1eaf-0a2f-4cf2-d68c-16a0b70bef92"
      },
      "source": [
        "model = tf.keras.models.load_model(\"odin_rn_model\")\n",
        "print(f\"Pre-trained model loaded with {model.count_params()/1e6} M parameters.\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Pre-trained model loaded with 0.572223 M parameters.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jhTEsfS58vvT"
      },
      "source": [
        "## Load the CIFAR-10 dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bbMcXWvw679h",
        "outputId": "f1e32f85-da6c-460d-824d-5ad6df334201"
      },
      "source": [
        "(_, _), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
        "print(f\"Total test examples: {len(x_test)}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Total test examples: 10000\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qjSdsdKl84hF"
      },
      "source": [
        "## Define constants"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T5_ZADvL7CSh"
      },
      "source": [
        "BATCH_SIZE = 128\n",
        "AUTO = tf.data.AUTOTUNE\n",
        "SAMPLES_PERTURB = 1000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WjSqTtJV86Ht"
      },
      "source": [
        "## Prepare data loaders"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TsWzYdFy7Y4k"
      },
      "source": [
        "perturb_samples = x_test[:SAMPLES_PERTURB].astype(\"float32\")\n",
        "perturb_ds = tf.data.Dataset.from_tensor_slices(perturb_samples).batch(BATCH_SIZE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9tMNGgs99It2"
      },
      "source": [
        "## Calculating the perturbation magnitude ($\\epsilon^{*}$)\n",
        "\n",
        "From the paper: \n",
        "\n",
        "> In our method, we search for the $\\epsilon^{*}$ which maximizes the score $S(x)$ with only the in-distribution validation dataset $D_{\\text {in }}^{\\text {val }}$:\n",
        "\n",
        "$$\\epsilon^{*}=\\underset{\\varepsilon}{\\arg \\max } \\sum_{x \\in D_{i n}^{\\text {val }}} S({\\boldsymbol{x}})$$\n",
        "\n",
        "$S(x)$ is given by:\n",
        "\n",
        ">  For out-of-distribution detection, we use the scoring function:\n",
        "<center>\n",
        "$S_{DeConf}(\\boldsymbol{x})=\\max _{i} h_{i}(\\boldsymbol{x})$ or $g(\\boldsymbol{x})$\n",
        "</center>\n",
        "\n",
        "Perturbation of an input image is realized using the equation below:\n",
        "\n",
        "$$\n",
        "\\hat{\\boldsymbol{x}}=\\boldsymbol{x}-\\epsilon \\operatorname{sign}\\left(-\\nabla_{\\boldsymbol{x}} S(\\boldsymbol{x})\\right)\n",
        "$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WzUie-5-8jAZ"
      },
      "source": [
        "# Let's define our model to obtain scores.\n",
        "scorer = tf.keras.Model(model.input, model.layers[-3].output)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JIboPp8bCa3k"
      },
      "source": [
        "# Grid as defined in Section 3.2.\n",
        "epsilon_grid = [0.0025, 0.005, 0.01, 0.02, 0.04, 0.08]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Si8J39lcDSnK"
      },
      "source": [
        "def perturb_images(model, epsilon):\n",
        "    batch_wise_means = []\n",
        "    \n",
        "    for images in perturb_ds:\n",
        "        test_ds_var = tf.Variable(images, trainable=True)\n",
        "        \n",
        "        with tf.GradientTape() as tape:\n",
        "            # Calculate the scores.\n",
        "            tape.watch(test_ds_var)\n",
        "            logits = model(test_ds_var, training=False)\n",
        "            loss = tf.reduce_max(logits, axis=1)\n",
        "            loss = -tf.reduce_mean(loss)\n",
        "\n",
        "        # Calculate the gradients of the scores with respect to the inputs.\n",
        "        gradients = tape.gradient(loss, test_ds_var)\n",
        "        gradients = tf.math.greater_equal(gradients, 0)\n",
        "        gradients = tf.cast(gradients, tf.float32)\n",
        "        gradients = (gradients - 0.5) * 2\n",
        "\n",
        "        # Perturb the inputs and derive new mean score.\n",
        "        # test_ds_var.assign_add(epsilon * gradients)\n",
        "        static_tensor = tf.convert_to_tensor(test_ds_var)\n",
        "        static_tensor = static_tensor - epsilon * gradients\n",
        "        static_tensor = tf.clip_by_value(static_tensor, 0., 255.)\n",
        "        \n",
        "        new_scores = model.predict(static_tensor)\n",
        "        new_scores = -tf.reduce_max(new_scores, axis=1)\n",
        "        new_mean_score = tf.reduce_mean(new_scores).numpy()\n",
        "        batch_wise_means.append(new_mean_score)\n",
        "    \n",
        "    return batch_wise_means"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d44O5i5HFwjF",
        "outputId": "97283960-b882-4d4e-d33c-2700cafe1f1f"
      },
      "source": [
        "# Derive the perturbation magnitude. \n",
        "mean_scores = {}\n",
        "\n",
        "for epsilon in epsilon_grid:\n",
        "    mean_scores[epsilon] = np.mean(perturb_images(scorer, epsilon))\n",
        "\n",
        "best_epsilon = min(mean_scores, key=(lambda key: mean_scores[key]))\n",
        "print(f\"Epsilon: {best_epsilon / 2.}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epsilon: 0.04\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}