site/en/tutorials/understanding/sngp.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9teObmxrP0FE"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mz8tfSwOP4fW"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CBwpoARnQ3H-"
      },
      "source": [
        "# Uncertainty-aware Deep Learning with SNGP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2dL6_obQRBGQ"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/tutorials/understanding/sngp\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/understanding/sngp.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/tutorials/understanding/sngp.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/tutorials/understanding/sngp.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UvW1QEMtP7Gy"
      },
      "source": [
        "In AI applications that are safety-critical, such as medical decision making and autonomous driving, or where the data is inherently noisy (for example, natural language understanding), it is important for a deep classifier to reliably quantify its uncertainty. The deep classifier should be able to be aware of its own limitations and when it should hand control over to the human experts. This tutorial shows how to improve a deep classifier's ability in quantifying uncertainty using a technique called **Spectral-normalized Neural Gaussian Process ([SNGP](https://arxiv.org/abs/2006.10108){.external})**.\n",
        "\n",
        "The core idea of SNGP is to improve a deep classifier's _**distance awareness**_ by applying simple modifications to the network. A model's _distance awareness_ is a measure of how its predictive probability reflects the distance between the test example and the training data. This is a desirable property that is common for gold-standard probabilistic models (for example, the [Gaussian process](https://en.wikipedia.org/wiki/Gaussian_process){.external} with RBF kernels) but is lacking in models with deep neural networks. SNGP provides a simple way to inject this Gaussian-process behavior into a deep classifier while maintaining its predictive accuracy.\n",
        "\n",
        "This tutorial implements a deep residual network (ResNet)-based SNGP model on [scikit-learn’s two moons](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_moons.html){.external} dataset, and compares its uncertainty surface with that of two other popular uncertainty approaches: [Monte Carlo dropout](https://arxiv.org/abs/1506.02142){.external} and [Deep ensemble](https://arxiv.org/abs/1612.01474){.external}.\n",
        "\n",
        "This tutorial illustrates the SNGP model on a toy 2D dataset. For an example of applying SNGP to a real-world natural language understanding task using a BERT-base, check out the [SNGP-BERT tutorial](https://www.tensorflow.org/text/tutorials/uncertainty_quantification_with_sngp_bert). For high-quality implementations of an SNGP model (and many other uncertainty methods) on a wide variety of benchmark datasets (such as [CIFAR-100](https://www.tensorflow.org/datasets/catalog/cifar100), [ImageNet](https://www.tensorflow.org/datasets/catalog/imagenet2012), [Jigsaw toxicity detection](https://www.tensorflow.org/datasets/catalog/wikipedia_toxicity_subtypes), etc), refer to the [Uncertainty Baselines](https://github.com/google/uncertainty-baselines){.external} benchmark."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K-tZzAIlvMv_"
      },
      "source": [
        "## About SNGP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ysyslHCyvYi-"
      },
      "source": [
        "SNGP is a simple approach to improve a deep classifier's uncertainty quality while maintaining a similar level of accuracy and latency. Given a deep residual network, SNGP makes two simple changes to the model:\n",
        "\n",
        "* It applies spectral normalization to the hidden residual layers.\n",
        "* It replaces the Dense output layer with a Gaussian process layer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ffU8rJ_eWHou"
      },
      "source": [
        ">![SNGP](http://tensorflow.org/tutorials/understanding/images/sngp.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2L88PoKr6XaE"
      },
      "source": [
        "Compared to other uncertainty approaches (such as Monte Carlo dropout or Deep ensemble), SNGP has several advantages:\n",
        "\n",
        "* It works for a wide range of state-of-the-art residual-based architectures (for example, (Wide) ResNet, DenseNet, or BERT).\n",
        "* It is a single-model method—it does not rely on ensemble averaging). Therefore, SNGP has a similar level of latency as a single deterministic network, and can be scaled easily to large datasets like [ImageNet](https://github.com/google/uncertainty-baselines/tree/main/baselines/imagenet){.external} and [Jigsaw Toxic Comments classification](https://github.com/google/uncertainty-baselines/tree/main/baselines/toxic_comments){.external}.\n",
        "* It has strong out-of-domain detection performance due to the _distance-awareness_ property.\n",
        "\n",
        "The downsides of this method are:\n",
        "\n",
        "* The predictive uncertainty of SNGP is computed using the [Laplace approximation](http://www.gaussianprocess.org/gpml/chapters/RW3.pdf){.external}. Therefore, theoretically, the posterior uncertainty of SNGP is different from that of an exact Gaussian process.\n",
        "\n",
        "* SNGP training needs a covariance reset step at the beginning of a new epoch. This can add a tiny amount of extra complexity to a training pipeline. This tutorial shows a simple way to implement this using Keras callbacks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ck_O7S8r1boS"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOS9qFlW2o3J"
      },
      "outputs": [],
      "source": [
        "!pip install -U -q --use-deprecated=legacy-resolver tf-models-official tensorflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RCoGSYz-PIR7"
      },
      "outputs": [],
      "source": [
        "# refresh pkg_resources so it takes the changes into account.\n",
        "import pkg_resources\n",
        "import importlib\n",
        "importlib.reload(pkg_resources)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJKtMbtYNXHn"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.colors as colors\n",
        "\n",
        "import sklearn.datasets\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "import official.nlp.modeling.layers as nlp_layers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8RCyTrQO4ZXo"
      },
      "source": [
        "Define visualization macros"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o5VTrcxo3rI-"
      },
      "outputs": [],
      "source": [
        "plt.rcParams['figure.dpi'] = 140\n",
        "\n",
        "DEFAULT_X_RANGE = (-3.5, 3.5)\n",
        "DEFAULT_Y_RANGE = (-2.5, 2.5)\n",
        "DEFAULT_CMAP = colors.ListedColormap([\"#377eb8\", \"#ff7f00\"])\n",
        "DEFAULT_NORM = colors.Normalize(vmin=0, vmax=1,)\n",
        "DEFAULT_N_GRID = 100"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lL5HzdYT5x5J"
      },
      "source": [
        "## The two moon dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nqazrSzhd24R"
      },
      "source": [
        "Create the training and evaluation datasets from the [scikit-learn two moon dataset](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_moons.html){.external}."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xJ_ul9Ii52mh"
      },
      "outputs": [],
      "source": [
        "def make_training_data(sample_size=500):\n",
        "  \"\"\"Create two moon training dataset.\"\"\"\n",
        "  train_examples, train_labels = sklearn.datasets.make_moons(\n",
        "      n_samples=2 * sample_size, noise=0.1)\n",
        "\n",
        "  # Adjust data position slightly.\n",
        "  train_examples[train_labels == 0] += [-0.1, 0.2]\n",
        "  train_examples[train_labels == 1] += [0.1, -0.2]\n",
        "\n",
        "  return train_examples, train_labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "goQkrxR_fFGd"
      },
      "source": [
        "Evaluate the model's predictive behavior over the entire 2D input space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cj7ldkNw5-cT"
      },
      "outputs": [],
      "source": [
        "def make_testing_data(x_range=DEFAULT_X_RANGE, y_range=DEFAULT_Y_RANGE, n_grid=DEFAULT_N_GRID):\n",
        "  \"\"\"Create a mesh grid in 2D space.\"\"\"\n",
        "  # testing data (mesh grid over data space)\n",
        "  x = np.linspace(x_range[0], x_range[1], n_grid)\n",
        "  y = np.linspace(y_range[0], y_range[1], n_grid)\n",
        "  xv, yv = np.meshgrid(x, y)\n",
        "  return np.stack([xv.flatten(), yv.flatten()], axis=-1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G9BYe4yqfeFa"
      },
      "source": [
        "To evaluate model uncertainty, add an out-of-domain (OOD) dataset that belongs to a third class. The model never observes these OOD examples during training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5UHz2SU4feSI"
      },
      "outputs": [],
      "source": [
        "def make_ood_data(sample_size=500, means=(2.5, -1.75), vars=(0.01, 0.01)):\n",
        "  return np.random.multivariate_normal(\n",
        "      means, cov=np.diag(vars), size=sample_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l1-jmJb45_at"
      },
      "outputs": [],
      "source": [
        "# Load the train, test and OOD datasets.\n",
        "train_examples, train_labels = make_training_data(\n",
        "    sample_size=500)\n",
        "test_examples = make_testing_data()\n",
        "ood_examples = make_ood_data(sample_size=500)\n",
        "\n",
        "# Visualize\n",
        "pos_examples = train_examples[train_labels == 0]\n",
        "neg_examples = train_examples[train_labels == 1]\n",
        "\n",
        "plt.figure(figsize=(7, 5.5))\n",
        "\n",
        "plt.scatter(pos_examples[:, 0], pos_examples[:, 1], c=\"#377eb8\", alpha=0.5)\n",
        "plt.scatter(neg_examples[:, 0], neg_examples[:, 1], c=\"#ff7f00\", alpha=0.5)\n",
        "plt.scatter(ood_examples[:, 0], ood_examples[:, 1], c=\"red\", alpha=0.1)\n",
        "\n",
        "plt.legend([\"Positive\", \"Negative\", \"Out-of-Domain\"])\n",
        "\n",
        "plt.ylim(DEFAULT_Y_RANGE)\n",
        "plt.xlim(DEFAULT_X_RANGE)\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nlzxsnBBkybB"
      },
      "source": [
        "Here, the blue and orange represent the positive and negative classes, and the red represents the OOD data. A model that quantifies the uncertainty well is expected to be confident when close to training data  (i.e., $p(x_{test})$ close to 0 or 1), and be uncertain when far away from the training data regions  (i.e., $p(x_{test})$ close to 0.5)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RJ3i4n8li-Mv"
      },
      "source": [
        "## The deterministic model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Utncgxs3lc4u"
      },
      "source": [
        "### Define model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-z83-nq6jPlb"
      },
      "source": [
        "Start from the (baseline) deterministic model: a multi-layer residual network (ResNet) with dropout regularization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wCBRm8fc6CgY"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "class DeepResNet(tf.keras.Model):\n",
        "  \"\"\"Defines a multi-layer residual network.\"\"\"\n",
        "  def __init__(self, num_classes, num_layers=3, num_hidden=128,\n",
        "               dropout_rate=0.1, **classifier_kwargs):\n",
        "    super().__init__()\n",
        "    # Defines class meta data.\n",
        "    self.num_hidden = num_hidden\n",
        "    self.num_layers = num_layers\n",
        "    self.dropout_rate = dropout_rate\n",
        "    self.classifier_kwargs = classifier_kwargs\n",
        "\n",
        "    # Defines the hidden layers.\n",
        "    self.input_layer = tf.keras.layers.Dense(self.num_hidden, trainable=False)\n",
        "    self.dense_layers = [self.make_dense_layer() for _ in range(num_layers)]\n",
        "\n",
        "    # Defines the output layer.\n",
        "    self.classifier = self.make_output_layer(num_classes)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    # Projects the 2d input data to high dimension.\n",
        "    hidden = self.input_layer(inputs)\n",
        "\n",
        "    # Computes the ResNet hidden representations.\n",
        "    for i in range(self.num_layers):\n",
        "      resid = self.dense_layers[i](hidden)\n",
        "      resid = tf.keras.layers.Dropout(self.dropout_rate)(resid)\n",
        "      hidden += resid\n",
        "\n",
        "    return self.classifier(hidden)\n",
        "\n",
        "  def make_dense_layer(self):\n",
        "    \"\"\"Uses the Dense layer as the hidden layer.\"\"\"\n",
        "    return tf.keras.layers.Dense(self.num_hidden, activation=\"relu\")\n",
        "\n",
        "  def make_output_layer(self, num_classes):\n",
        "    \"\"\"Uses the Dense layer as the output layer.\"\"\"\n",
        "    return tf.keras.layers.Dense(\n",
        "        num_classes, **self.classifier_kwargs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4u870GAen2aO"
      },
      "source": [
        "This tutorial uses a six-layer ResNet with 128 hidden units."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bWL9wCnGpc4h"
      },
      "outputs": [],
      "source": [
        "resnet_config = dict(num_classes=2, num_layers=6, num_hidden=128)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I47RY26wurgg"
      },
      "outputs": [],
      "source": [
        "resnet_model = DeepResNet(**resnet_config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "okQXf2F1ur16"
      },
      "outputs": [],
      "source": [
        "resnet_model.build((None, 2))\n",
        "resnet_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ZzueLoImW0t"
      },
      "source": [
        "### Train model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JfQ1k-IUukLt"
      },
      "source": [
        "Configure the training parameters to use `SparseCategoricalCrossentropy` as the loss function and the Adam optimizer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a9ZkfNXAnumV"
      },
      "outputs": [],
      "source": [
        "loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "metrics = tf.keras.metrics.SparseCategoricalAccuracy(),\n",
        "optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=1e-4)\n",
        "\n",
        "train_config = dict(loss=loss, metrics=metrics, optimizer=optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wEWzfHvf6A5_"
      },
      "source": [
        "Train the model for 100 epochs with batch size 128."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Fx6EtGXpzVr"
      },
      "outputs": [],
      "source": [
        "fit_config = dict(batch_size=128, epochs=100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cFwkbKIuqj7Y"
      },
      "outputs": [],
      "source": [
        "resnet_model.compile(**train_config)\n",
        "resnet_model.fit(train_examples, train_labels, **fit_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eo6tKKd1rvBh"
      },
      "source": [
        "### Visualize uncertainty"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HZDMX7gZrZ-5"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "def plot_uncertainty_surface(test_uncertainty, ax, cmap=None):\n",
        "  \"\"\"Visualizes the 2D uncertainty surface.\n",
        "  \n",
        "  For simplicity, assume these objects already exist in the memory:\n",
        "\n",
        "    test_examples: Array of test examples, shape (num_test, 2).\n",
        "    train_labels: Array of train labels, shape (num_train, ).\n",
        "    train_examples: Array of train examples, shape (num_train, 2).\n",
        "  \n",
        "  Arguments:\n",
        "    test_uncertainty: Array of uncertainty scores, shape (num_test,).\n",
        "    ax: A matplotlib Axes object that specifies a matplotlib figure.\n",
        "    cmap: A matplotlib colormap object specifying the palette of the\n",
        "      predictive surface.\n",
        "\n",
        "  Returns:\n",
        "    pcm: A matplotlib PathCollection object that contains the palette\n",
        "      information of the uncertainty plot.\n",
        "  \"\"\"\n",
        "  # Normalize uncertainty for better visualization.\n",
        "  test_uncertainty = test_uncertainty / np.max(test_uncertainty)\n",
        "\n",
        "  # Set view limits.\n",
        "  ax.set_ylim(DEFAULT_Y_RANGE)\n",
        "  ax.set_xlim(DEFAULT_X_RANGE)\n",
        "\n",
        "  # Plot normalized uncertainty surface.\n",
        "  pcm = ax.imshow(\n",
        "      np.reshape(test_uncertainty, [DEFAULT_N_GRID, DEFAULT_N_GRID]),\n",
        "      cmap=cmap,\n",
        "      origin=\"lower\",\n",
        "      extent=DEFAULT_X_RANGE + DEFAULT_Y_RANGE,\n",
        "      vmin=DEFAULT_NORM.vmin,\n",
        "      vmax=DEFAULT_NORM.vmax,\n",
        "      interpolation='bicubic',\n",
        "      aspect='auto')\n",
        "\n",
        "  # Plot training data.\n",
        "  ax.scatter(train_examples[:, 0], train_examples[:, 1],\n",
        "             c=train_labels, cmap=DEFAULT_CMAP, alpha=0.5)\n",
        "  ax.scatter(ood_examples[:, 0], ood_examples[:, 1], c=\"red\", alpha=0.1)\n",
        "\n",
        "  return pcm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a1age2y0339T"
      },
      "source": [
        "Now visualize the predictions of the deterministic model. First plot the class probability:\n",
        "$$p(x) = softmax(logit(x))$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4aqFgQOD40lb"
      },
      "outputs": [],
      "source": [
        "resnet_logits = resnet_model(test_examples)\n",
        "resnet_probs = tf.nn.softmax(resnet_logits, axis=-1)[:, 0]  # Take the probability for class 0."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rpflH2Qj33oN"
      },
      "outputs": [],
      "source": [
        "_, ax = plt.subplots(figsize=(7, 5.5))\n",
        "\n",
        "pcm = plot_uncertainty_surface(resnet_probs, ax=ax)\n",
        "\n",
        "plt.colorbar(pcm, ax=ax)\n",
        "plt.title(\"Class Probability, Deterministic Model\")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ShGAB7FNYgU"
      },
      "source": [
        "In this plot, the yellow and purple are the predictive probabilities for the two classes. The deterministic model did a good job in classifying the two known classes—blue and orange—with a nonlinear decision boundary. However, it is not **distance-aware**, and classified the never-observed red out-of-domain (OOD) examples confidently as the orange class.\n",
        "\n",
        "Visualize the model uncertainty by computing the [predictive variance](https://en.wikipedia.org/wiki/Bernoulli_distribution#Variance):\n",
        "$$var(x) = p(x) * (1 - p(x))$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VxMce5D15Qwt"
      },
      "outputs": [],
      "source": [
        "resnet_uncertainty = resnet_probs * (1 - resnet_probs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fwdEfb5_6woh"
      },
      "outputs": [],
      "source": [
        "_, ax = plt.subplots(figsize=(7, 5.5))\n",
        "\n",
        "pcm = plot_uncertainty_surface(resnet_uncertainty, ax=ax)\n",
        "\n",
        "plt.colorbar(pcm, ax=ax)\n",
        "plt.title(\"Predictive Uncertainty, Deterministic Model\")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alVunBEXhD-1"
      },
      "source": [
        "In this plot, the yellow indicates high uncertainty, and the purple indicates low uncertainty. A deterministic ResNet's uncertainty depends only on the test examples' distance from the decision boundary. This leads the model to be over-confident when out of the training domain. The next section shows how SNGP behaves differently on this dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CwDa80iOh32J"
      },
      "source": [
        "## The SNGP model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IbjNrmLu8oXv"
      },
      "source": [
        "### Define SNGP model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "urtuCrk1Hf5P"
      },
      "source": [
        "Let's now implement the SNGP model. Both the SNGP components, `SpectralNormalization` and `RandomFeatureGaussianProcess`, are available at the tensorflow_model's [built-in  layers](https://github.com/tensorflow/models/tree/master/official/nlp/modeling/layers). "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7IlQwKCEGwpk"
      },
      "source": [
        ">![SNGP](http://tensorflow.org/tutorials/understanding/images/sngp.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rp2O2iv8LSke"
      },
      "source": [
        "Let's inspect these two components in more detail. (You can also jump to [the full SNGP model](#full-sngp-model) section to learn how SNGP is implemented.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5n4NIt3QjKwl"
      },
      "source": [
        "#### `SpectralNormalization` wrapper"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tE-Va7J2jR2X"
      },
      "source": [
        "[`SpectralNormalization`](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/layers/spectral_normalization.py){.external} is a Keras layer wrapper. It can be applied to an existing Dense layer like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dp8vqJBWLSq3"
      },
      "outputs": [],
      "source": [
        "dense = tf.keras.layers.Dense(units=10)\n",
        "dense = nlp_layers.SpectralNormalization(dense, norm_multiplier=0.9)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9q25_6fRJRh"
      },
      "source": [
        "Spectral normalization regularizes the hidden weight $W$ by gradually guiding its spectral norm (that is, the largest eigenvalue of $W$) toward the target value `norm_multiplier`).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cqt-DVzvqyAE"
      },
      "source": [
        "Note: Usually it is preferable to set `norm_multiplier` to a value smaller than 1. However in practice, it can be also relaxed to a larger value to ensure the deep network has enough expressive power."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yqvxJUXBjBhV"
      },
      "source": [
        "#### The Gaussian Process (GP) layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7rYfIgtrjHnB"
      },
      "source": [
        "[`RandomFeatureGaussianProcess`](https://github.com/tensorflow/models/blob/master/official/nlp/modeling/layers/gaussian_process.py){.external} implements a [random-feature based approximation](https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf){.external} to a Gaussian process model that is end-to-end trainable with a deep neural network. Under the hood, the Gaussian process layer implements a two-layer network:\n",
        "\n",
        "$$logits(x) = \\Phi(x) \\beta, \\quad \\Phi(x)=\\sqrt{\\frac{2}{M}} * cos(Wx + b)$$\n",
        "\n",
        "Here, $x$ is the input, and $W$ and $b$ are frozen weights initialized randomly from Gaussian and Uniform distributions, respectively. (Therefore, $\\Phi(x)$ are called \"random features\".) $\\beta$ is the learnable kernel weight similar to that of a Dense layer. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NqnU39ui3wAE"
      },
      "outputs": [],
      "source": [
        "batch_size = 32\n",
        "input_dim = 1024\n",
        "num_classes = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LrlVd-foRJno"
      },
      "outputs": [],
      "source": [
        "gp_layer = nlp_layers.RandomFeatureGaussianProcess(units=num_classes,\n",
        "                                               num_inducing=1024,\n",
        "                                               normalize_input=False,\n",
        "                                               scale_random_features=True,\n",
        "                                               gp_cov_momentum=-1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xxb8sSAg5AGf"
      },
      "source": [
        "The main parameters of the GP layers are:\n",
        "\n",
        "* `units`: The dimension of the output logits.\n",
        "* `num_inducing`: The dimension $M$ of the hidden weight $W$. Default to 1024.\n",
        "* `normalize_input`: Whether to apply layer normalization to the input $x$.\n",
        "* `scale_random_features`: Whether to apply the scale $\\sqrt{2/M}$ to the hidden output.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgzw09gS03ae"
      },
      "source": [
        "Note: For a deep neural network that is sensitive to the learning rate (for example, ResNet-50 and ResNet-110), it is generally recommended to set `normalize_input=True` to stabilize training, and set `scale_random_features=False` to avoid the learning rate from being modified in unexpected ways when passing through the GP layer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pZkcKw-u7XRp"
      },
      "source": [
        "* `gp_cov_momentum` controls how the model covariance is computed. If set to a positive value (for example, `0.999`), the covariance matrix is computed using the momentum-based moving average update (similar to batch normalization). If set to `-1`, the covariance matrix is updated without momentum."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P13X7Adt-c2d"
      },
      "source": [
        "Note: The momentum-based update method can be sensitive to batch size. Therefore it is generally recommended to set  `gp_cov_momentum=-1` to compute the covariance exactly. For this to work properly, the covariance matrix estimator needs to be reset at the beginning of a new epoch in order to avoid counting the same data twice. For `RandomFeatureGaussianProcess`, this can be done by calling its `reset_covariance_matrix()`. The next section shows an easy implementation of this using Keras' built-in API.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0AA492qA1biZ"
      },
      "source": [
        "Given a batch input with shape `(batch_size, input_dim)`, the GP layer returns a `logits` tensor (shape `(batch_size, num_classes)`) for prediction, and also  `covmat` tensor (shape `(batch_size, batch_size)`) which is the posterior covariance matrix of the batch logits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yOXxBYnMi1v4"
      },
      "outputs": [],
      "source": [
        "embedding = tf.random.normal(shape=(batch_size, input_dim))\n",
        "\n",
        "logits, covmat = gp_layer(embedding)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ALBqcAtwDNiO"
      },
      "source": [
        "Note: Notice that under this implementation of the SNGP model, the predictive logits $logit(x_{test})$ for all classes share the same covariance matrix $var(x_{test})$, which describes the distance between $x_{test}$ from the training data.\n",
        "\n",
        "Theoretically, it is possible to extend the algorithm to compute different variance values for different classes (as introduced in the [original SNGP paper](https://arxiv.org/abs/2006.10108){.external}). However, this is difficult to scale to problems with large output spaces (such as classification with ImageNet or language modeling)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "II3GVzjJhu5Z"
      },
      "source": [
        "<a name=\"full-sngp-model\"></a>\n",
        "#### The full SNGP model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F5Fm0NUlLTHd"
      },
      "source": [
        "Given the base class `DeepResNet`, the SNGP model can be implemented easily by  modifying the residual network's hidden and output layers. For compatibility with Keras `model.fit()` API, also modify the model's `call()` method so it only outputs `logits` during training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dzx4FsO97QZv"
      },
      "outputs": [],
      "source": [
        "class DeepResNetSNGP(DeepResNet):\n",
        "  def __init__(self, spec_norm_bound=0.9, **kwargs):\n",
        "    self.spec_norm_bound = spec_norm_bound\n",
        "    super().__init__(**kwargs)\n",
        "\n",
        "  def make_dense_layer(self):\n",
        "    \"\"\"Applies spectral normalization to the hidden layer.\"\"\"\n",
        "    dense_layer = super().make_dense_layer()\n",
        "    return nlp_layers.SpectralNormalization(\n",
        "        dense_layer, norm_multiplier=self.spec_norm_bound)\n",
        "\n",
        "  def make_output_layer(self, num_classes):\n",
        "    \"\"\"Uses Gaussian process as the output layer.\"\"\"\n",
        "    return nlp_layers.RandomFeatureGaussianProcess(\n",
        "        num_classes,\n",
        "        gp_cov_momentum=-1,\n",
        "        **self.classifier_kwargs)\n",
        "\n",
        "  def call(self, inputs, training=False, return_covmat=False):\n",
        "    # Gets logits and a covariance matrix from the GP layer.\n",
        "    logits, covmat = super().call(inputs)\n",
        "\n",
        "    # Returns only logits during training.\n",
        "    if not training and return_covmat:\n",
        "      return logits, covmat\n",
        "\n",
        "    return logits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02SvlnDFyszm"
      },
      "source": [
        "Use the same architecture as the deterministic model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QiDcC5ipyqMU"
      },
      "outputs": [],
      "source": [
        "resnet_config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D9-imQtDyZL5"
      },
      "outputs": [],
      "source": [
        "sngp_model = DeepResNetSNGP(**resnet_config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "12P9UxlCyq6k"
      },
      "outputs": [],
      "source": [
        "sngp_model.build((None, 2))\n",
        "sngp_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BlCgMRrwkXgN"
      },
      "source": [
        "<a name=\"covariance-reset-callback\"></a>\n",
        "Implement a Keras callback to reset the covariance matrix at the beginning of a new epoch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W6Wr8D_0n-cQ"
      },
      "outputs": [],
      "source": [
        "class ResetCovarianceCallback(tf.keras.callbacks.Callback):\n",
        "\n",
        "  def on_epoch_begin(self, epoch, logs=None):\n",
        "    \"\"\"Resets covariance matrix at the beginning of the epoch.\"\"\"\n",
        "    if epoch > 0:\n",
        "      self.model.classifier.reset_covariance_matrix()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMiz_mYawEPh"
      },
      "source": [
        "Add this callback to the `DeepResNetSNGP` model class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yrSi0ZD5zaDf"
      },
      "outputs": [],
      "source": [
        "class DeepResNetSNGPWithCovReset(DeepResNetSNGP):\n",
        "  def fit(self, *args, **kwargs):\n",
        "    \"\"\"Adds ResetCovarianceCallback to model callbacks.\"\"\"\n",
        "    kwargs[\"callbacks\"] = list(kwargs.get(\"callbacks\", []))\n",
        "    kwargs[\"callbacks\"].append(ResetCovarianceCallback())\n",
        "\n",
        "    return super().fit(*args, **kwargs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "asIwYqGlwJcP"
      },
      "source": [
        "### Train model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4YRzayOCopt9"
      },
      "source": [
        "Use `tf.keras.model.fit` to train the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "coazo53nwJqv"
      },
      "outputs": [],
      "source": [
        "sngp_model = DeepResNetSNGPWithCovReset(**resnet_config)\n",
        "sngp_model.compile(**train_config)\n",
        "sngp_model.fit(train_examples, train_labels, **fit_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KTONd8vowgEP"
      },
      "source": [
        "### Visualize uncertainty"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yhkpsiy10l9d"
      },
      "source": [
        "First compute the predictive logits and  variances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bqRPqeEavi4Z"
      },
      "outputs": [],
      "source": [
        "sngp_logits, sngp_covmat = sngp_model(test_examples, return_covmat=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p7w0iW_L0cU8"
      },
      "outputs": [],
      "source": [
        "sngp_variance = tf.linalg.diag_part(sngp_covmat)[:, None]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YLbz_EZF1Gay"
      },
      "source": [
        "<a name=\"mean-field-logits\"></a>\n",
        "Now compute the posterior predictive probability. The classic method for computing the predictive probability of a probabilistic model is to use Monte Carlo sampling, i.e.,\n",
        "\n",
        "$$E(p(x)) = \\frac{1}{M} \\sum_{m=1}^M logit_m(x), $$\n",
        "\n",
        "where $M$ is the sample size, and $logit_m(x)$ are random samples from the SNGP posterior $MultivariateNormal$(`sngp_logits`,`sngp_covmat`). However, this approach can be slow for latency-sensitive applications such as autonomous driving or real-time bidding. Instead, you can approximate $E(p(x))$ using the [mean-field method](https://arxiv.org/abs/2006.07584){.external}:\n",
        "\n",
        "$$E(p(x)) \\approx softmax(\\frac{logit(x)}{\\sqrt{1+ \\lambda * \\sigma^2(x)}})$$\n",
        "\n",
        "where $\\sigma^2(x)$ is the SNGP variance, and $\\lambda$ is often chosen as $\\pi/8$ or $3/\\pi^2$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1A9NYMhd0iZ-"
      },
      "outputs": [],
      "source": [
        "sngp_logits_adjusted = sngp_logits / tf.sqrt(1. + (np.pi / 8.) * sngp_variance)\n",
        "sngp_probs = tf.nn.softmax(sngp_logits_adjusted, axis=-1)[:, 0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bNVs_KO-5HdL"
      },
      "source": [
        "Note: Instead of fixing $\\lambda$ to a fixed value, you can also treat it as a hyperparameter, and tune it to optimize the model's calibration performance. This is known as [temperature scaling](http://proceedings.mlr.press/v70/guo17a.html){.external} in the deep learning uncertainty literature. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DlYPlUbJfBFa"
      },
      "source": [
        "This mean-field method is implemented as a built-in function `layers.gaussian_process.mean_field_logits`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hgb3WSaY8iQY"
      },
      "outputs": [],
      "source": [
        "def compute_posterior_mean_probability(logits, covmat, lambda_param=np.pi / 8.):\n",
        "  # Computes uncertainty-adjusted logits using the built-in method.\n",
        "  logits_adjusted = nlp_layers.gaussian_process.mean_field_logits(\n",
        "      logits, covmat, mean_field_factor=lambda_param)\n",
        "  \n",
        "  return tf.nn.softmax(logits_adjusted, axis=-1)[:, 0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kVToZpG7QqS3"
      },
      "outputs": [],
      "source": [
        "sngp_logits, sngp_covmat = sngp_model(test_examples, return_covmat=True)\n",
        "sngp_probs = compute_posterior_mean_probability(sngp_logits, sngp_covmat)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bVi_Whpwe3O4"
      },
      "source": [
        "### SNGP Summary"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dRmmxuO41BV4"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "\n",
        "def plot_predictions(pred_probs, model_name=\"\"):\n",
        "  \"\"\"Plot normalized class probabilities and predictive uncertainties.\"\"\"\n",
        "  # Compute predictive uncertainty.\n",
        "  uncertainty = pred_probs * (1. - pred_probs)\n",
        "\n",
        "  # Initialize the plot axes.\n",
        "  fig, axs = plt.subplots(1, 2, figsize=(14, 5))\n",
        "\n",
        "  # Plots the class probability.\n",
        "  pcm_0 = plot_uncertainty_surface(pred_probs, ax=axs[0])\n",
        "  # Plots the predictive uncertainty.\n",
        "  pcm_1 = plot_uncertainty_surface(uncertainty, ax=axs[1])\n",
        "\n",
        "  # Adds color bars and titles.\n",
        "  fig.colorbar(pcm_0, ax=axs[0])\n",
        "  fig.colorbar(pcm_1, ax=axs[1])\n",
        "\n",
        "  axs[0].set_title(f\"Class Probability, {model_name}\")\n",
        "  axs[1].set_title(f\"(Normalized) Predictive Uncertainty, {model_name}\")\n",
        "\n",
        "  plt.show()  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R9kY5dJg8fEi"
      },
      "source": [
        "You can now put everything together. The entire procedure—training, evaluation and uncertainty computation—can be done in just five lines:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NQtUAG4-ftqe"
      },
      "outputs": [],
      "source": [
        "def train_and_test_sngp(train_examples, test_examples):\n",
        "  sngp_model = DeepResNetSNGPWithCovReset(**resnet_config)\n",
        "\n",
        "  sngp_model.compile(**train_config)\n",
        "  sngp_model.fit(train_examples, train_labels, verbose=0, **fit_config)\n",
        "\n",
        "  sngp_logits, sngp_covmat = sngp_model(test_examples, return_covmat=True)\n",
        "  sngp_probs = compute_posterior_mean_probability(sngp_logits, sngp_covmat)\n",
        "\n",
        "  return sngp_probs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dl3N7kHJ283w"
      },
      "outputs": [],
      "source": [
        "sngp_probs = train_and_test_sngp(train_examples, test_examples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OAUGrXSv6k3R"
      },
      "source": [
        "Visualize the class probability (left) and the predictive uncertainty (right) of the SNGP model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rxt3CY51A_jq"
      },
      "outputs": [],
      "source": [
        "plot_predictions(sngp_probs, model_name=\"SNGP\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "raHP5Vuiuku9"
      },
      "source": [
        "Remember that in the class probability plot (left), the yellow and purple are class probabilities. When close to the training data domain, SNGP correctly classifies the examples with high confidence (i.e., assigning near 0 or 1 probability). When far away from the training data, SNGP gradually becomes less confident, and its predictive probability becomes close to 0.5 while the (normalized) model uncertainty rises to 1."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4_VRi9w7Km3"
      },
      "source": [
        "Compare this to the uncertainty surface of the deterministic model: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aMAdstYZ7T_w"
      },
      "outputs": [],
      "source": [
        "plot_predictions(resnet_probs, model_name=\"Deterministic\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mao9L-LYE1Nl"
      },
      "source": [
        "As mentioned earlier, a deterministic model is not _distance-aware_. Its uncertainty is defined by the distance of the test example from the decision boundary. This leads the model to produce overconfident predictions for the out-of-domain examples (red)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EKURpzOf0oNq"
      },
      "source": [
        "## Comparison with other uncertainty approaches"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S1DPELWE6LL8"
      },
      "source": [
        "This section compares the uncertainty of SNGP with [Monte Carlo dropout](https://arxiv.org/abs/1506.02142){.external} and [Deep ensemble](https://arxiv.org/abs/1612.01474){.external}.\n",
        "\n",
        "Both of these methods are based on Monte Carlo averaging of multiple forward passes of deterministic models. First, set the ensemble size $M$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xLqkihbk8Dey"
      },
      "outputs": [],
      "source": [
        "num_ensemble = 10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jM5AQmtIAatd"
      },
      "source": [
        "### Monte Carlo dropout"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZBzp2LBt7-kj"
      },
      "source": [
        "Given a trained neural network with Dropout layers, Monte Carlo dropout computes the mean predictive probability\n",
        "\n",
        "$$E(p(x)) = \\frac{1}{M}\\sum_{m=1}^M softmax(logit_m(x))$$\n",
        "\n",
        "by averaging over multiple Dropout-enabled forward passes $\\{logit_m(x)\\}_{m=1}^M$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I7R2WBgq4-OC"
      },
      "outputs": [],
      "source": [
        "def mc_dropout_sampling(test_examples):\n",
        "  # Enable dropout during inference.\n",
        "  return resnet_model(test_examples, training=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r6oXgaDZAiD0"
      },
      "outputs": [],
      "source": [
        "# Monte Carlo dropout inference.\n",
        "dropout_logit_samples = [mc_dropout_sampling(test_examples) for _ in range(num_ensemble)]\n",
        "dropout_prob_samples = [tf.nn.softmax(dropout_logits, axis=-1)[:, 0] for dropout_logits in dropout_logit_samples]\n",
        "dropout_probs = tf.reduce_mean(dropout_prob_samples, axis=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_oUNtbVG-YuI"
      },
      "outputs": [],
      "source": [
        "dropout_probs = tf.reduce_mean(dropout_prob_samples, axis=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J-mhyp8hAiPn"
      },
      "outputs": [],
      "source": [
        "plot_predictions(dropout_probs, model_name=\"MC Dropout\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mwtj2vTB75cF"
      },
      "source": [
        "### Deep ensemble"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-Z2veGJ9ZgY"
      },
      "source": [
        "Deep ensemble is a state-of-the-art (but expensive) method for deep learning uncertainty. To train a Deep ensemble, first train $M$ ensemble members."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a43hxiJC3Kla"
      },
      "outputs": [],
      "source": [
        "# Deep ensemble training\n",
        "resnet_ensemble = []\n",
        "for _ in range(num_ensemble):\n",
        "  resnet_model = DeepResNet(**resnet_config)\n",
        "  resnet_model.compile(optimizer=optimizer, loss=loss, metrics=metrics)\n",
        "  resnet_model.fit(train_examples, train_labels, verbose=0, **fit_config)\n",
        "\n",
        "  resnet_ensemble.append(resnet_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Al7uM-fn_ZE1"
      },
      "source": [
        "Collect logits and compute the mean predictive probability $E(p(x)) = \\frac{1}{M}\\sum_{m=1}^M softmax(logit_m(x))$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c6E9PntV3Mue"
      },
      "outputs": [],
      "source": [
        "# Deep ensemble inference\n",
        "ensemble_logit_samples = [model(test_examples) for model in resnet_ensemble]\n",
        "ensemble_prob_samples = [tf.nn.softmax(logits, axis=-1)[:, 0] for logits in ensemble_logit_samples]\n",
        "ensemble_probs = tf.reduce_mean(ensemble_prob_samples, axis=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z_JhYftG-NKR"
      },
      "outputs": [],
      "source": [
        "plot_predictions(ensemble_probs, model_name=\"Deep ensemble\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GH33oVvV5-ez"
      },
      "source": [
        "Both the Monte Carlo Dropout and Deep ensemble methods improve the model's uncertainty ability by making the decision boundary less certain. However, they both inherit the deterministic deep network's limitation in lacking distance awareness."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x9nAaWuYfD03"
      },
      "source": [
        "## Summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ryMkRllFHOC"
      },
      "source": [
        "In this tutorial, you have:\n",
        "* Implemented the SNGP model on a deep classifier to improve its distance awareness.\n",
        "* Trained the SNGP model end-to-end using Keras `Model.fit` API.\n",
        "* Visualized the uncertainty behavior of SNGP.\n",
        "* Compared the uncertainty behavior between SNGP, Monte Carlo dropout and deep ensemble models."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UoTekiQmkZXF"
      },
      "source": [
        "## Resources and further reading"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HoIikRybke-b"
      },
      "source": [
        "* Check out the [SNGP-BERT tutorial](https://www.tensorflow.org/text/tutorials/uncertainty_quantification_with_sngp_bert) for an example of applying SNGP on a BERT model for uncertainty-aware natural language understanding.\n",
        "* Go to the [Uncertainty Baselines GitHub repo](https://github.com/google/uncertainty-baselines){.external} for the implementation of SNGP model (and many other uncertainty methods) on a wide variety of benchmark datasets (for example, [CIFAR](https://www.tensorflow.org/datasets/catalog/cifar100), [ImageNet](https://www.tensorflow.org/datasets/catalog/imagenet2012), [Jigsaw toxicity detection](https://www.tensorflow.org/datasets/catalog/wikipedia_toxicity_subtypes), etc).\n",
        "* For a deeper understanding of the SNGP method, check out the paper titled [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://arxiv.org/abs/2006.10108){.external}.\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "sngp.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}