{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BNN_supervised_airfoils.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "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.6.6"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s1i4P0Z-LFnl"
      },
      "source": [
        "# Bayesian Neural Nets for airfoil flows\n",
        "\n",
        "## Overview\n",
        "\n",
        "We are looking at the same setup as in the notebook $\\textit{Supervised training for RANS flows around airfoils}$: A turbulent airflow around wing profiles, for which we'd like to know the average motion\n",
        "and pressure distribution around this airfoil for different Reynolds numbers and angles of attack. In the other notebook, we tackled this by completely bypassing any physical solver and instead training a neural network that learns the quantities of interest. Now, we want to extend this approach to so-called Bayesian Neural Networks (short: BNNs). In contrast to traditional networks, that learn a single point estimate for each weight value, BNNs aim at learning a $\\textit{distribution}$ over each weight parameter (e.g. a Gaussian with mean $\\mu$ and variance $\\sigma^{2}$). During a forward-pass, each parameter in the network is then sampled from its corresponding distribution. In that sense, the network parameters themselves are $\\textit{random variables}$ and each forward pass becomes $\\textit{stochastic}$, because for a given input the predictions will vary with every forward-pass. This allows to assess how $\\textit{uncertain}$ the network is: If the predictions vary a lot, we think that the network is uncertain about its output. \n",
        "\n",
        "\n",
        "## A quick ride through the theory behind BNNs\n",
        "Instead of fixed point estimates of the parameters, we aim at learning distributions now. When dealing with BNNs, one therefore assumes a prior distribution $p({w})$ over the network weights ${w}$. Typically, independent Gaussians are used, but other distributions are possible, too. When observing data ${D}$, the prior distribution over the weights is updated to the so-called posterior according to Bayes rule:\n",
        "\n",
        "$\\begin{equation}\n",
        "p({w}|{D})=\\frac{p({D}|{w})p({w})}{p({D})}\n",
        "\\end{equation}\n",
        "$.\n",
        "\n",
        "This computation is usually intractable, especially when dealing with non-trivial network architectures. Therefore, the posterior $p({w}|{D})$ is approximated with an easy-to-evaluate variational distribution $q_{\\theta}(w)$, parametrized by $\\theta$. Again, independent Normal distributions are typically used for each weight. The optimization goal is then to find a distribution $q_{\\theta}(w)$ that is close to the true posterior. This can be achieved by minimizing the so-called evidence lower bound (ELBO): \n",
        "\n",
        "$\\begin{equation}\n",
        "\\mathcal{L}(\\theta)=  -E_{q_{\\theta}}[\\log(p(D|{w}))] +KL(q_{\\theta}({w})||p({w}))\n",
        "\\end{equation}\n",
        "$, \n",
        "\n",
        "which is the loss function for BNNs. The first term is an expected negative log-likelhood of the data. Minimizing it means explaining the data as well as possible. In practice, this first term is often an conventional loss functions like mean squared error (note that MSE can be seen as negative log-likelihood for normal noise with unit variance). The second term is the KL-divergence between the approximate posterior and the prior. Minimizing it means encouraging the network weight distributions to stay close to the prior distribution. In that sense, the terms have opposite goals: The first time encourages the model to explain the data as well as possible, whereas the second term encourages the model to stay close to the (random) prior distributions, which implies randomness and regularization.\n",
        "\n",
        " In practice, the expectation of the log-likelihood is approximated by a single sample during a forward pass. This is, in every forward pass the network weights are sampled from the current posterior distribution, the loss is computed, and the error is backpropagated in order to update the posterior distribution parameters. If you are curious on how one can backpropagate through distributions, you can read about it [here](https://arxiv.org/abs/1505.05424). An excellent, thorough introduction to Bayesian Neural Networks is available in Yarin Gals [thesis](https://mlg.eng.cam.ac.uk/yarin/thesis/thesis.pdf) in chapters 2 & 3.\n",
        "\n",
        "\n",
        "But now let's jump right in and discover the BNN features along the way!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jzE1wwAjG_QQ"
      },
      "source": [
        "### Read in Data\n",
        "Like in the previous notebook we'll skip the data generation process. This example is adapted from [this codebase](https://github.com/thunil/Deep-Flow-Prediction), which you can check out for details. Here, we'll simply download a small set of training data generated with a Spalart-Almaras RANS simulation in [OpenFOAM](https://openfoam.org/)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JwZudtWauiGa"
      },
      "source": [
        "# uncomment for final notebook version\n",
        "# import numpy as np\n",
        "# import os.path, random\n",
        "\n",
        "# if not os.path.isfile('data-airfoils.npz'):\n",
        "#     import urllib.request\n",
        "#     url=\"https://ge.in.tum.de/download/2019-deepFlowPred/data.npz\"\n",
        "#     print(\"Downloading training data (300MB), this can take a few minutes the first time...\")\n",
        "#     urllib.request.urlretrieve(url, 'data-airfoils.npz')\n",
        "\n",
        "# npfile=np.load('data-airfoils.npz')\n",
        "# print(\"Loaded data, {} training, {} validation samples\".format(len(npfile[\"inputs\"]),len(npfile[\"vinputs\"])))\n",
        "\n",
        "# print(\"Size of the inputs array: \"+format(npfile[\"inputs\"].shape))\n",
        "\n",
        "# # reshape to channels_last for convencience\n",
        "# X_train = np.moveaxis(npfile[\"inputs\"],1,-1)\n",
        "# y_train = np.moveaxis(npfile[\"targets\"],1,-1)\n",
        "\n",
        "# X_val = np.moveaxis(npfile[\"vinputs\"],1,-1)\n",
        "# y_val = np.moveaxis(npfile[\"vtargets\"],1,-1)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HaocUpT5DGgC",
        "outputId": "7469de2a-d484-48fd-a2d7-2c38ead8c64f"
      },
      "source": [
        "# Use mount for convenience - will disappear in final notebook\n",
        "import numpy as np\n",
        "import os.path, random\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "npfile=np.load('gdrive/My Drive/data.npz')\n",
        "\n",
        "\n",
        "print(\"Loaded data, {} training, {} validation samples\".format(len(npfile[\"inputs\"]),len(npfile[\"vinputs\"])))\n",
        "\n",
        "print(\"Size of the inputs array: \"+format(npfile[\"inputs\"].shape))\n",
        "\n",
        "# reshape to channels_last for convencience\n",
        "X_train = np.moveaxis(npfile[\"inputs\"],1,-1)\n",
        "y_train = np.moveaxis(npfile[\"targets\"],1,-1)\n",
        "\n",
        "X_val = np.moveaxis(npfile[\"vinputs\"],1,-1)\n",
        "y_val = np.moveaxis(npfile[\"vtargets\"],1,-1)\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/gdrive; to attempt to forcibly remount, call drive.mount(\"/content/gdrive\", force_remount=True).\n",
            "Loaded data, 320 training, 80 validation samples\n",
            "Size of the inputs array: (320, 3, 128, 128)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5C_Ooq0JG_Qk"
      },
      "source": [
        "### Look at Data\n",
        "Now we have some training data. We can look at it using the code we also used in the original notebook:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1y1zHmKlAWNs",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 327
        },
        "outputId": "3d71bd09-13aa-4a13-99f7-e69c8d167710"
      },
      "source": [
        "from PIL import Image\n",
        "from matplotlib import cm\n",
        "\n",
        "def show(a, stats=False): # helper to show three target channels, normalized with colormap\n",
        "    for i in range(3):\n",
        "        b = np.flipud(a[i].transpose())\n",
        "        min,mean,max = np.min(b),np.mean(b),np.max(b); \n",
        "        if stats: print([b[0,0],min,mean,max])\n",
        "        b -= min; b /= (max-min)\n",
        "        if stats: print(\"{}\".format([\"p:\",\"ux:\",\"uy:\"][i])) # note, for inputs the first channel is the mask...\n",
        "        display(Image.fromarray( cm.magma(b, bytes=True) ))\n",
        "\n",
        "def showSbs(a1,a2, stats=False): # helper to show three target channels, normalized with colormap side by side\n",
        "    c=[]\n",
        "    for i in range(3):\n",
        "        b = np.flipud( np.concatenate((a2[:,:,i],a1[:,:,i]),axis=1).transpose())\n",
        "        min, mean, max = np.min(b), np.mean(b), np.max(b); \n",
        "        if stats: print([b[0,0],min,mean,max])\n",
        "        b -= min; b /= (max-min)\n",
        "        c.append(b)\n",
        "    c = np.concatenate(c,axis=1)\n",
        "    display(Image.fromarray( cm.magma(c, bytes=True) ))\n",
        "\n",
        "num=40\n",
        "print(\"\\nHere are all 3 inputs are shown at the top (mask,in x, in y) \\nSide by side with the 3 output channels (p,vx,vy) at the bottom:\")\n",
        "showSbs(X_train[num],y_train[num])\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Here are all 3 inputs are shown at the top (mask,in x, in y) \n",
            "Side by side with the 3 output channels (p,vx,vy) at the bottom:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 33032 characters]\n",
            "text/plain": [
              "<PIL.Image.Image image mode=RGBA size=384x256 at 0x7FC2D5F00390>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2gdKINAG_Qs"
      },
      "source": [
        "### Model Definition\n",
        "Now let's look at how we can implement BNNs. Instead of PyTorch, we will use TensorFlow, in particular the extension TensorFlow Probability, which has easy-to-implement probabilistic layers. Like in the other notebook, we want to implement a U-Net, consisting of Convolutional blocks with skip-layer connections. For now, we only want to make the decoder, i.e. second part of the U-Net bayesian. For this, we will take advantage of TensorFlows flipout layers (in particular, the convolutional implementation). In a forward pass, those layers automatically sample from the current posterior distribution and store the KL-divergence between prior and posterior in _model.losses_. One can specify the desired divergence measure (typically KL-divergence) and modify the prior and approximate posterior distributions, if other than normal distributions are desired. Other than that, the flipout layers can be used just like regular layers in sequential models. The code below implements a single convolutional block of the U-Net:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9-8O6qJ-G_Qt"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_probability.python.distributions as tfd\n",
        "from tensorflow.keras import Sequential\n",
        "from tensorflow.keras.initializers import RandomNormal\n",
        "from tensorflow.keras.layers import Input, Conv2D, Conv2DTranspose,UpSampling2D, BatchNormalization, ReLU, LeakyReLU, SpatialDropout2D, MaxPooling2D\n",
        "from tensorflow_probability.python.layers import Convolution2DFlipout\n",
        "from tensorflow.keras.models import Model\n",
        "\n",
        "def tfBlockUnet(filters=3,transposed=False, kernel_size=3,bn=True,relu=True, pad=\"same\", dropout=0., flipout=False,\n",
        "              kdf=None,name=''):\n",
        "    block = Sequential(name=name)\n",
        "    if relu:\n",
        "        block.add(ReLU())\n",
        "    else:\n",
        "        block.add(LeakyReLU(0.2))\n",
        "    if not transposed:\n",
        "        block.add(Conv2D(filters=filters, kernel_size=kernel_size, padding=pad,\n",
        "                         kernel_initializer=RandomNormal(0.0, 0.02), activation=None,strides=(2,2)))\n",
        "    else:\n",
        "        block.add(UpSampling2D(interpolation = 'bilinear'))\n",
        "        if flipout:\n",
        "            block.add(Convolution2DFlipout(filters=filters, kernel_size=(kernel_size-1), strides=(1, 1), padding=pad,\n",
        "                                           data_format=\"channels_last\", kernel_divergence_fn=kdf,\n",
        "                                           activation=None))\n",
        "        else:\n",
        "            block.add(Conv2D(filters=filters, kernel_size=(kernel_size-1), padding=pad,\n",
        "                         kernel_initializer=RandomNormal(0.0, 0.02), strides=(1,1), activation=None))\n",
        "    \n",
        "    block.add(SpatialDropout2D(rate=dropout))\n",
        "    \n",
        "    if bn:\n",
        "        block.add(BatchNormalization(axis=-1, epsilon=1e-05,momentum=0.9))\n",
        "\n",
        "    return block"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "57AMSMp5sPFZ"
      },
      "source": [
        "And now let's use those blocks to define the full network - the structure is almost idenetical to the other notebook. We are now manually defining the kernel-divergence function as kdf and rescaling it with a factor called _kl_scaling_. There are two reasons for this: \n",
        "\n",
        "First, we should only apply the kl-divergence once per epoch if we want to use the correct loss (like introduced on the top of this notebook). Since we will use batch training, we should rescale the Kl-divergence by the number of batches, such that in every parameter update only _kl-div / num_batches_ is added to the loss. After one epoch, _num_batches_ parameter updates were performed and the 'full' KL-divergence was used. \n",
        "\n",
        "Second, by scaling the KL-divergence part of the loss up or down, we have a way of tuning how much randomness we want to allow in the network: If we neglect the KL-divergence completely, we would just minimize the regular loss (e.g. MSE or MAE), like in a conventional neural network. If we instead neglect the negative-log-likelihood, we would optimize the network such that we obtain random draws from the prior distribution. Balancing those extremes can be done by fine-tuning the scaling of the KL-divergence and is hard in practice. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pGoGJwqdG_Qv"
      },
      "source": [
        "def Bayes_DfpNet(input_shape=(128,128,3),expo=5,dropout=0.,flipout=False,kl_scaling=10000):\n",
        "    channels = int(2 ** expo + 0.5)\n",
        "    kdf = (lambda q, p, _: tfd.kl_divergence(q, p) / tf.cast(kl_scaling, dtype=tf.float32))\n",
        "\n",
        "    layer1=Sequential(name='layer1')\n",
        "    layer1.add(Conv2D(filters=channels,kernel_size=4,strides=(2,2),padding='same',activation=None,data_format='channels_last'))\n",
        "    layer2=tfBlockUnet(filters=channels*2,transposed=False,bn=True, relu=False,dropout=dropout,name='layer2')\n",
        "    layer3=tfBlockUnet(filters=channels*2,transposed=False,bn=True, relu=False,dropout=dropout,name='layer3')        \n",
        "    layer4=tfBlockUnet(filters=channels*4,transposed=False,bn=True, relu=False,dropout=dropout,name='layer4')        \n",
        "    layer5=tfBlockUnet(filters=channels*8,transposed=False,bn=True, relu=False,dropout=dropout,name='layer5')        \n",
        "    layer6=tfBlockUnet(filters=channels*8,transposed=False,bn=True, relu=False,dropout=dropout,kernel_size=2,pad='valid',name='layer6')        \n",
        "    layer7=tfBlockUnet(filters=channels*8,transposed=False,bn=True, relu=False,dropout=dropout,kernel_size=2,pad='valid',name='layer7')        \n",
        "\n",
        "    # note, kernel size is internally reduced by one for the decoder part\n",
        "    dlayer7=tfBlockUnet(filters=channels*8,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf, kernel_size=2,pad='valid',name='dlayer7')        \n",
        "    dlayer6=tfBlockUnet(filters=channels*8,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,kernel_size=2,pad='valid',name='dlayer6')        \n",
        "    dlayer5=tfBlockUnet(filters=channels*4,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer5')        \n",
        "    dlayer4=tfBlockUnet(filters=channels*2,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer4')\n",
        "    dlayer3=tfBlockUnet(filters=channels*2,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer3')\n",
        "    dlayer2=tfBlockUnet(filters=channels,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer2')\n",
        "    dlayer1=Sequential(name='outlayer')\n",
        "    dlayer1.add(ReLU())\n",
        "    dlayer1.add(Conv2DTranspose(3,kernel_size=4,strides=(2,2),padding='same'))\n",
        "\n",
        "    # forward pass\n",
        "    inputs=Input(input_shape)\n",
        "    out1 = layer1(inputs)\n",
        "    out2 = layer2(out1)\n",
        "    out3 = layer3(out2)\n",
        "    out4 = layer4(out3)\n",
        "    out5 = layer5(out4)\n",
        "    out6 = layer6(out5)\n",
        "    out7 = layer7(out6)\n",
        "    # ... bottleneck ...\n",
        "    dout6 = dlayer7(out7)\n",
        "    dout6_out6 = tf.concat([dout6,out6],axis=3)\n",
        "    dout6 = dlayer6(dout6_out6)\n",
        "    dout6_out5 = tf.concat([dout6, out5], axis=3)\n",
        "    dout5 = dlayer5(dout6_out5)\n",
        "    dout5_out4 = tf.concat([dout5, out4], axis=3)\n",
        "    dout4 = dlayer4(dout5_out4)\n",
        "    dout4_out3 = tf.concat([dout4, out3], axis=3)\n",
        "    dout3 = dlayer3(dout4_out3)\n",
        "    dout3_out2 = tf.concat([dout3, out2], axis=3)\n",
        "    dout2 = dlayer2(dout3_out2)\n",
        "    dout2_out1 = tf.concat([dout2, out1], axis=3)\n",
        "    dout1 = dlayer1(dout2_out1)\n",
        "    \n",
        "    return Model(inputs=inputs,outputs=dout1)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ETzjWZJzyMEu"
      },
      "source": [
        "Let's define the hyperparameters and create a tensorflow dataset to organize inputs and targets. Since we have 320 observations in the training set, for a batch-size of 10 we should rescale the KL-divergence with a factor of 320/10=32 in order apply the full KL-divergence just once per epoch. We will further scale the KL-divergence down by another factor of 5000, which has shown to work well in practice. \n",
        "\n",
        "Furthermore, we will define a function that implements learning rate decay. Intuitvely, this allows the optimization to be more precise (by making smaller steps) in later epochs, while still making fast progress (by making bigger steps) in the first epochs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H23a0JAJclEN"
      },
      "source": [
        "import math \n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "batch_size=10\n",
        "lrG=0.001\n",
        "epochs = 150\n",
        "kl_pref = 5000\n",
        "\n",
        "dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(len(X_train),\n",
        "    seed=46168531, reshuffle_each_iteration=False).batch(batch_size, drop_remainder=False)\n",
        "\n",
        "def computeLR(i, epochs, minLR, maxLR):\n",
        "  if i < epochs * 0.5:\n",
        "      return maxLR\n",
        "  e = (i / float(epochs) - 0.5) * 2.\n",
        "  # rescale second half to min/max range\n",
        "  fmin = 0.\n",
        "  fmax = 6.\n",
        "  e = fmin + e * (fmax - fmin)\n",
        "  f = math.pow(0.5, e)\n",
        "  return minLR + (maxLR - minLR) * f"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3lUU7A0o1PzV"
      },
      "source": [
        "We can visualize the learning rate decay: We start off with a constant rate and after half of the epochs we start to decay it exponentially, until arriving at half of the original learning rate."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tDXZO0FN02do",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 297
        },
        "outputId": "f2e3dbb3-faac-464a-8da5-d0f2939faa27"
      },
      "source": [
        "lrs=[computeLR(i, epochs, 0.5*lrG,lrG) for i in range(epochs)]\n",
        "plt.plot(lrs)\n",
        "plt.xlabel('Iteration')\n",
        "plt.ylabel('Learning Rate')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Text(0, 0.5, 'Learning Rate')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 64
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 13636 characters]\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lcwUJ6iH1uTr"
      },
      "source": [
        "Let's initialize the net."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LGaH2ZW73M1b",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e39ddaee-41f7-427a-82bd-4dba385c650a"
      },
      "source": [
        "from tensorflow.keras.optimizers import RMSprop, Adam\n",
        "\n",
        "model=Bayes_DfpNet(expo=3,flipout=True,kl_scaling=kl_pref*len(X_train)/batch_size)\n",
        "optimizer = Adam(learning_rate=lrG, beta_1=0.5,beta_2=0.9999)\n",
        "\n",
        "num_params = np.sum([np.prod(v.get_shape().as_list()) for v in model.trainable_variables])\n",
        "print('Our Bayesian U-Net has {} parameters.'.format(num_params))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:2281: UserWarning: `layer.add_variable` is deprecated and will be removed in a future version. Please use `layer.add_weight` method instead.\n",
            "  warnings.warn('`layer.add_variable` is deprecated and '\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Our Bayesian U-Net has 133475 parameters.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eF4UTqeB6Y4u"
      },
      "source": [
        "Due to some differences in the implementation, the number of parameters is slightly lower than in the conventional neural network (147555), even though we use the same `expo` parameter. In general, flipout layers however come with twice as many parameters compared to their conventional counterparts, since instead of a single point estimate one has to learn both mean and variance parameters (for a Gaussian posterior).\n",
        "\n",
        "Now we are ready to train!\n",
        "\n",
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VEQuKBegcd1K",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ce3a6b9c-8c14-4bee-be69-acb42740330c"
      },
      "source": [
        "from tensorflow.keras.losses import mae\n",
        "import math\n",
        "\n",
        "kl_losses=[]\n",
        "mae_losses=[]\n",
        "total_losses=[]\n",
        "val_maes=[]\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    # compute learning rate - decay is implemented\n",
        "    currLr = computeLR(epoch,epochs,0.5*lrG,lrG)\n",
        "    if currLr < lrG:\n",
        "            tf.keras.backend.set_value(optimizer.lr, currLr)\n",
        "\n",
        "    # iterate through training data \n",
        "    kl_sum = 0\n",
        "    mae_sum = 0\n",
        "    total_sum=0\n",
        "    for i, traindata in enumerate(dataset, 0):\n",
        "      # forward pass and loss computation\n",
        "        with tf.GradientTape() as tape:\n",
        "            inputs, targets = traindata\n",
        "            prediction = model(inputs)\n",
        "            loss_mae = tf.reduce_mean(mae(prediction, targets))\n",
        "            kl=sum(model.losses)\n",
        "            loss_value=kl+tf.cast(loss_mae, dtype='float32')\n",
        "        # backpropagate gradients and update parameters \n",
        "        gradients = tape.gradient(loss_value, model.trainable_variables)\n",
        "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "   \n",
        "        # store losses per batch\n",
        "        kl_sum += kl\n",
        "        mae_sum += tf.reduce_mean(loss_mae)\n",
        "        total_sum+=tf.reduce_mean(loss_value)\n",
        "\n",
        "    # store losses per epoch\n",
        "    kl_losses+=[kl_sum/len(dataset)]\n",
        "    mae_losses+=[mae_sum/len(dataset)]\n",
        "    total_losses+=[total_sum/len(dataset)]\n",
        "\n",
        "    # validation\n",
        "    outputs = model.predict(X_val)\n",
        "    val_maes += [tf.reduce_mean(mae(y_val, outputs))]\n",
        "\n",
        "    if epoch<3 or epoch%20==0:\n",
        "        print('Finnished epoch {}...'.format(epoch))\n",
        "        print('total loss: {}'.format(total_losses[-1]))\n",
        "        print('KL loss: {}'.format(kl_losses[-1]))\n",
        "        print('MAE loss: {}'.format(mae_losses[-1]))\n",
        "        print('Validation MAE loss: {}'.format(val_maes[-1]))\n",
        "        print('----------------------------')\n",
        "\n",
        "    "
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Finnished epoch 0...\n",
            "total loss: 0.6994723081588745\n",
            "KL loss: 0.5650137066841125\n",
            "MAE loss: 0.13445863978867173\n",
            "Validation MAE loss: 0.08071838319301605\n",
            "----------------------------\n",
            "Finnished epoch 1...\n",
            "total loss: 0.629783034324646\n",
            "KL loss: 0.5600312352180481\n",
            "MAE loss: 0.06975189820671622\n",
            "Validation MAE loss: 0.05924928933382034\n",
            "----------------------------\n",
            "Finnished epoch 2...\n",
            "total loss: 0.6130387187004089\n",
            "KL loss: 0.5546823143959045\n",
            "MAE loss: 0.05835640023620712\n",
            "Validation MAE loss: 0.053308360278606415\n",
            "----------------------------\n",
            "Finnished epoch 20...\n",
            "total loss: 0.48715242743492126\n",
            "KL loss: 0.45129039883613586\n",
            "MAE loss: 0.035862083967624137\n",
            "Validation MAE loss: 0.03532093018293381\n",
            "----------------------------\n",
            "Finnished epoch 40...\n",
            "total loss: 0.3690536916255951\n",
            "KL loss: 0.33777734637260437\n",
            "MAE loss: 0.03127635233989842\n",
            "Validation MAE loss: 0.02977834641933441\n",
            "----------------------------\n",
            "Finnished epoch 60...\n",
            "total loss: 0.26104623079299927\n",
            "KL loss: 0.23509857058525085\n",
            "MAE loss: 0.025947713498049033\n",
            "Validation MAE loss: 0.025638097897171974\n",
            "----------------------------\n",
            "Finnished epoch 80...\n",
            "total loss: 0.17463460564613342\n",
            "KL loss: 0.15130069851875305\n",
            "MAE loss: 0.02333393547606524\n",
            "Validation MAE loss: 0.022063611075282097\n",
            "----------------------------\n",
            "Finnished epoch 100...\n",
            "total loss: 0.1251290738582611\n",
            "KL loss: 0.10345939546823502\n",
            "MAE loss: 0.021669673259241923\n",
            "Validation MAE loss: 0.020878834649920464\n",
            "----------------------------\n",
            "Finnished epoch 120...\n",
            "total loss: 0.09566079825162888\n",
            "KL loss: 0.07473748922348022\n",
            "MAE loss: 0.020923318253538117\n",
            "Validation MAE loss: 0.020196828991174698\n",
            "----------------------------\n",
            "Finnished epoch 140...\n",
            "total loss: 0.07551293075084686\n",
            "KL loss: 0.055130843073129654\n",
            "MAE loss: 0.020382081854780156\n",
            "Validation MAE loss: 0.020393211394548416\n",
            "----------------------------\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7aM5Ra2C7k1v"
      },
      "source": [
        "The model is trained! Let's look at the loss. Since the loss consists of two separate parts, it is helpful to monitor both parts (MAE and KL)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3niTOL6CcrHo",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 299
        },
        "outputId": "4c77fabf-09b4-447e-cab1-7f6ea584762e"
      },
      "source": [
        "fig,axs=plt.subplots(ncols=3,nrows=1,figsize=(20,4))\n",
        "axs[0].plot(kl_losses,color='red')\n",
        "axs[0].set_title('KL Loss (Train)')\n",
        "axs[1].plot(mae_losses,color='blue',label='train')\n",
        "axs[1].plot(val_maes,color='green',label='val')\n",
        "axs[1].legend()\n",
        "axs[1].set_title('MAE Loss')\n",
        "axs[2].plot(total_losses,label='Total',color='black')\n",
        "axs[2].plot(kl_losses,label='KL',color='red')\n",
        "axs[2].plot(mae_losses,label='MAE',color='blue')\n",
        "axs[2].set_title('Total Train Loss')\n",
        "axs[2].legend()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.legend.Legend at 0x7fc2cdd7bd10>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 67
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 42640 characters]\n",
            "text/plain": [
              "<Figure size 1440x288 with 3 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WBGaJng6psg-"
      },
      "source": [
        "This way, we can double-check if minimizing one part of the loss comes at the cost of increasing the other. For our case, we observe that both parts decrease smoothly. In particular the MAE loss is not increasing for the validation set, indicating that we are not overfitting. \n",
        "\n",
        "It is good practice to double-check how many layers added KL-losses. We can inspect _model.losses_ for that. Since the decoder consists of 6 sequential blocks with flipout layers, we expect 6 entries in _model.losses_."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LusTjxWRdRjn",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ebd12958-c36b-42b8-b9ea-1b622727f7fb"
      },
      "source": [
        "# there should be 6 entries in model.losses since we have 6 blocks with flipout layers in our model\n",
        "print('There are {} entries in model.losses'.format(len(model.losses)))\n",
        "print(model.losses)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "There are 6 entries in model.losses\n",
            "[<tf.Tensor: shape=(), dtype=float32, numpy=0.0023392083>, <tf.Tensor: shape=(), dtype=float32, numpy=0.00630205>, <tf.Tensor: shape=(), dtype=float32, numpy=0.01660155>, <tf.Tensor: shape=(), dtype=float32, numpy=0.0093718935>, <tf.Tensor: shape=(), dtype=float32, numpy=0.0072054924>, <tf.Tensor: shape=(), dtype=float32, numpy=0.006180173>]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "auFzYkx0G_Q6"
      },
      "source": [
        "Now let's visualize how the BNN performs for unseen data from the validation set. Since a forward pass is stochastic for BNNs, we will perform several forward passes for each input. This allows us to compute an average prediction (which is equivalent to marginalizing out the posterior distribution) and a standard deviation for each pixel."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z-gyY4LKG_Q7"
      },
      "source": [
        "reps=20\n",
        "preds=np.zeros(shape=(reps,)+X_val.shape)\n",
        "for rep in range(reps):\n",
        "    preds[rep,:,:,:,:]=model.predict(X_val)\n",
        "preds_mean=np.mean(preds,axis=0)\n",
        "preds_std=np.std(preds,axis=0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcSUWaiHcVkb"
      },
      "source": [
        "In the first plot, the input is shown in the first row. The second row illustrates the result of a single forward pass. \n",
        "\n",
        "In the second plot, the average prediction is shown in the first row, and the corresponding uncertainty in the second row. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 565
        },
        "id": "WSEXdGc7G_Q8",
        "outputId": "d3879aa3-b88d-436d-e607-6cfb054e1d13"
      },
      "source": [
        "num=16\n",
        "# show some single prediction\n",
        "print('First Row: Inputs, Second row: Single forward Pass')\n",
        "showSbs(y_val[num],preds[0][num])\n",
        "\n",
        "#show avg. pred in first row, std.deviation over prediction in second row\n",
        "print('First Row: Average Prediction, Second row: Std. Deviation')\n",
        "showSbs(preds_mean[num],preds_std[num])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "First Row: Inputs, Second row: Single forward Pass\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 84848 characters]\n",
            "text/plain": [
              "<PIL.Image.Image image mode=RGBA size=384x256 at 0x7FC2CDAC8210>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "First Row: Average Prediction, Second row: Std. Deviation\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 81904 characters]\n",
            "text/plain": [
              "<PIL.Image.Image image mode=RGBA size=384x256 at 0x7FC2E2A27E10>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_b9HoUskeU1B"
      },
      "source": [
        "We can visualize the uncertainty over predictions more clearly by inspecting several samples from the posterior distribution for a given input. Below is code for a function that visualizes precisely that."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 166
        },
        "id": "gAcKTBCGmohC",
        "outputId": "880b2976-73ec-4e4f-fc9b-aed0852f7580"
      },
      "source": [
        "# plot repeated samples from posterior for some observations\n",
        "def plot_BNN_predictions(target, preds, pred_mean, pred_std, num_preds=5,channel=0):\n",
        "  if num_preds>len(preds):\n",
        "    print('num_preds was set to {}, but has to be smaller than the length of preds. Setting it to {}'.format(num_preds,len(preds)))\n",
        "    num_preds = len(preds)\n",
        "\n",
        "  # transpose and concatenate the frames that are to plot\n",
        "  to_plot=np.concatenate((target[:,:,channel].transpose().reshape(128,128,1),preds[0:num_preds,:,:,channel].transpose(), \n",
        "                          pred_mean[:,:,channel].transpose().reshape(128,128,1),pred_std[:,:,channel].transpose().reshape(128,128,1)),axis=-1)\n",
        "  fig, axs = plt.subplots(nrows=1,ncols=to_plot.shape[-1],figsize=(20,4))\n",
        "  for i in range(to_plot.shape[-1]):\n",
        "    label='Target' if i==0 else ('Avg Pred' if i == (num_preds+1) else ('Std Dev (normalized)' if i == (num_preds+2) else 'Pred {}'.format(i)))\n",
        "    frame = np.flipud(to_plot[:,:,i])\n",
        "    min=np.min(frame); max = np.max(frame)\n",
        "    frame -= min; frame /=(max-min)\n",
        "    axs[i].imshow(frame)\n",
        "    axs[i].axis('off')\n",
        "    axs[i].set_title(label)\n",
        "\n",
        "obs_idx=5\n",
        "plot_BNN_predictions(y_val[obs_idx,...],preds[:,obs_idx,:,:,:],preds_mean[obs_idx,...],preds_std[obs_idx,...])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 150864 characters]\n",
            "text/plain": [
              "<Figure size 1440x288 with 8 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ln7sn-UJpHa7"
      },
      "source": [
        "We are looking at channel 0, i.e. the pressure here. One can observe that the blue and yellow regions vary quite a bit across predictions. It is reassuring to note that - at least from visual inspection - the average (i.e. marginal) prediction is closer to the target than most of the single forward passes. \n",
        "\n",
        "It should also be noted that each frame was normalized for the visualization. Therefore, when looking at the uncertainty frame, we can inferr where the network is uncertain, but now how uncertain it is in absolute values. \n",
        "\n",
        "In order to assess a global measure of uncertainty we can however compute an average standard deviation over all samples in the validation set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Ja9z5yU8RJ3",
        "outputId": "a19db741-baf7-4eeb-e720-d40fc62b7ef6"
      },
      "source": [
        "# Let's look at how uncertain the network is, i.e. average uncertainty over average prediction\n",
        "uncertainty = np.round(np.mean(np.abs(preds_std),axis=(0,1,2))/np.mean(np.abs(preds_mean),axis=(0,1,2)),2)\n",
        "print(\"\\nAverage relative uncertainty on validation set: \\n pressure: {}, \\n ux: {},\\n uy: {}\".format(uncertainty[0],uncertainty[1],uncertainty[2]))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Average uncertainty on validation set: \n",
            " pressure: 0.26, \n",
            " ux: 0.03,\n",
            " uy: 0.16\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rn8Pc8plZM1D",
        "outputId": "cc5c8329-89a7-4a31-c21a-156dea1e621e"
      },
      "source": [
        "# average pixel-wise uncertainty over average pixel-wise prediction\n",
        "uncertainty = np.round(np.mean(preds_std/np.abs(preds_mean),axis=(0,1,2)),2)\n",
        "print(\"\\nAverage relative pixe-wise uncertainty on validation set: \\n pressure: {}, \\n ux: {},\\n uy: {}\".format(uncertainty[0],uncertainty[1],uncertainty[2]))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Average relative pixe-wise uncertainty on validation set: \n",
            " pressure: 1.7, \n",
            " ux: 0.27,\n",
            " uy: 1.28\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hcVF9S-JoprL"
      },
      "source": [
        "## Test evaluation\n",
        "\n",
        "Like in the case for a conventional neural network, let's now look at **proper** test samples, i.e. OOD samples, for which in this case we'll use new airfoil shapes. These are shapes that the network never saw in any training samples, and hence it tells us a bit about how well the model generalizes to new shapes.\n",
        "\n",
        "As these samples are at least slightly OOD, we can draw conclusions about how well the model generalizes, which the validation data would not really tell us. In particular, we would like to investigate if the model is more uncertain when handling OOD data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ClH53HqtOw3L",
        "outputId": "0d07d103-b8a0-4ee4-c61d-52acefc38bcd"
      },
      "source": [
        "if not os.path.isfile('data-airfoils-test.npz'):\n",
        "  import urllib.request\n",
        "  url=\"https://ge.in.tum.de/download/2019-deepFlowPred/data_test.npz\"\n",
        "  print(\"Downloading test data, this should be fast...\")\n",
        "  urllib.request.urlretrieve(url, 'data-airfoils-test.npz')\n",
        "\n",
        "nptfile=np.load('data-airfoils-test.npz')\n",
        "print(\"Loaded {}/{} test samples\".format(len(nptfile[\"test_inputs\"]),len(nptfile[\"test_targets\"])))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loaded 10/10 test samples\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TP5D9rhkNlsf",
        "outputId": "b8a036ad-be26-4258-a206-b8eb3ff70a86"
      },
      "source": [
        "X_test = np.moveaxis(nptfile[\"test_inputs\"],1,-1)\n",
        "y_test = np.moveaxis(nptfile[\"test_targets\"],1,-1)\n",
        "\n",
        "reps=10\n",
        "preds_test=np.zeros(shape=(reps,)+X_test.shape)\n",
        "for rep in range(reps):\n",
        "    preds_test[rep,:,:,:,:]=model.predict(X_test)\n",
        "preds_test_mean=np.mean(preds_test,axis=0)\n",
        "preds_test_std=np.std(preds_test,axis=0)\n",
        "test_loss = tf.reduce_mean(mae(preds_test_mean, y_test))\n",
        "\n",
        "print(\"\\nAverage test error: {}\".format(test_loss))\n",
        "\n",
        "\n",
        "# Let's look at how uncertain the network is\n",
        "uncertainty_std = np.round(np.mean(np.abs(preds_test_std),axis=(0,1,2))/np.mean(np.abs(preds_test_mean),axis=(0,1,2)),2)\n",
        "print(\"\\nAverage relative uncertainty on test set: \\n pressure: {}, \\n ux: {},\\n uy: {}\".format(uncertainty_std[0],uncertainty_std[1],uncertainty_std[2]))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Average test error: 0.025036948244479572\n",
            "\n",
            "Average relative uncertainty on test set: \n",
            " pressure: 0.27, \n",
            " ux: 0.04,\n",
            " uy: 0.18\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "shMWpZ9Te4J9",
        "outputId": "05930702-f52b-4bae-a7df-ab8905474895"
      },
      "source": [
        "# average pixel-wise uncertainty over average pixel-wise prediction\n",
        "uncertainty = np.round(np.mean(preds_test_std/np.abs(preds_test_mean),axis=(0,1,2)),2)\n",
        "print(\"\\nAverage relative pixe-wise uncertainty on test set: \\n pressure: {}, \\n ux: {},\\n uy: {}\".format(uncertainty[0],uncertainty[1],uncertainty[2]))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Average relative pixe-wise uncertainty on test set: \n",
            " pressure: 3.24, \n",
            " ux: 0.18,\n",
            " uy: 1.58\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "icgfvAIqoMpE"
      },
      "source": [
        "## Check this, might not be true any more\n",
        "This is reassuring: The error on the OOD test set with new shapes is higher than on the validation set. However, also the uncertainty is larger.\n",
        "In general it is hard to obtain a calibrated uncertainty estimate, but since we are dealing with a fairly simple problem here, it seems that the BNN could estimate it reasonably well."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 297
        },
        "id": "x_Nzzc7rrJOc",
        "outputId": "d8d3b5d2-e25c-4627-840a-7f59757f7098"
      },
      "source": [
        "# Why not to use pixel-wise uncertainty: Estimate is dominated by few pixels with low average prediction\n",
        "# since we divide by pred_mean, a value of pred_mean close to 0 can lead to exploding relative uncertainties\n",
        "plt.scatter(preds_mean[:,:,:,0],(preds_std[:,:,:,0]/np.abs(preds_mean[:,:,:,0])).flatten())\n",
        "plt.xlabel('Avg. predicted Pixel Value')\n",
        "plt.ylabel('Relative Uncertainty')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Text(0, 0.5, 'Relative Uncertainty')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 120
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 11000 characters]\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 166
        },
        "id": "oP1HOciVjh-o",
        "outputId": "4e04731b-c3e5-4873-fc11-419086d37581"
      },
      "source": [
        "obs_idx=5\n",
        "plot_BNN_predictions(y_test[obs_idx,...],preds_test[:,obs_idx,:,:,:],preds_test_mean[obs_idx,...],preds_test_std[obs_idx,...])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 138120 characters]\n",
            "text/plain": [
              "<Figure size 1440x288 with 8 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "trAG3xoQG_Q-"
      },
      "source": [
        "Despite these promising results, there are still many issues with Bayesian Neural Nets, limiting their use in most practical applications. One serious drawback is the need for additional scaling of the KL-loss and the fact that there is no convincing argument on why it is necessary yet (read e.g. [here](http://proceedings.mlr.press/v119/wenzel20a/wenzel20a.pdf) or [here](https://arxiv.org/abs/2008.05912)).\n",
        "Furthermore, some people think that assuming independent normal distributions as variational approximations to the posterior is an oversimplification since in practice the weights are actually highly correlated ([paper](https://arxiv.org/abs/1909.00719)). Other people instead argue that this might not be an issue, as long as the networks in use are deep enough ([paper](https://arxiv.org/abs/2002.03704)). On top of that, there is research on different (e.g. heavy-tailed) priors other than normals and many other aspects of BNNs."
      ]
    }
  ]
}
