tutorials/04_probability/probabilistic_model_by_tfp_layers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/hellocybernetics/TensorFlow2.0_Eager_Execution_Tutorials/blob/master/tutorials/99_others/probabilistic_model_by_tfp_layers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PToRtUfKhkgX"
   },
   "outputs": [],
   "source": [
    "# !pip install -q --upgrade tf-nightly-2.0-preview\n",
    "# !pip install -q tfp-nightly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XLrt-DYZhsuB"
   },
   "source": [
    "## TFP layers basic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rZgkVwaw2TRn"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.style.use(\"seaborn\")\n",
    "\n",
    "tfk = tf.keras\n",
    "tfkl = tf.keras.layers\n",
    "tfd = tfp.distributions\n",
    "tfpl = tfp.layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dJHXvcW1onZk"
   },
   "source": [
    "### simple regression\n",
    "There are four parameters $w_\\mu, w_\\sigma, b_\\mu, b_\\sigma$ .\n",
    "\n",
    "$$\n",
    "p(y\\mid x, w_\\mu, w_\\sigma, b_\\mu, b_\\sigma) = {\\mathcal N} (y \\mid w_\\mu x + b_\\mu,  w_\\sigma x+b_\\sigma)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wR83I3Vlh2J2"
   },
   "outputs": [],
   "source": [
    "input_dim = 1\n",
    "output_dim = 1\n",
    "\n",
    "model = tfk.Sequential([\n",
    "    tfkl.InputLayer([input_dim]),\n",
    "    tfkl.Dense(tfpl.IndependentNormal.params_size(output_dim)),\n",
    "    tfpl.IndependentNormal(output_dim)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 233
    },
    "colab_type": "code",
    "id": "SshzZfaPoGl_",
    "outputId": "b4209638-0c2d-4a5f-a5ff-f2bb1590444f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense (Dense)                (None, 2)                 4         \n",
      "_________________________________________________________________\n",
      "independent_normal (Independ ((None, 1), (None, 1))    0         \n",
      "=================================================================\n",
      "Total params: 4\n",
      "Trainable params: 4\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VbyJCeW8h7kw"
   },
   "outputs": [],
   "source": [
    "X = tf.constant([[3.],\n",
    "                 [2.],\n",
    "                 [-5.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "hegJ0EM7nt57",
    "outputId": "697cd30b-6aa3-42c4-8119-8c770c2bb670"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.Independent 'sequential_independent_normal_IndependentNormal_Independentsequential_independent_normal_IndependentNormal_Normal' batch_shape=[3] event_shape=[1] dtype=float32>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 89
    },
    "colab_type": "code",
    "id": "2It2hTA_Eads",
    "outputId": "b90197e5-8cfe-4356-dfe5-06555ae2c8af"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=227, shape=(3, 1), dtype=float32, numpy=\n",
       "array([[-0.43018448],\n",
       "       [-0.30774704],\n",
       "       [-7.049762  ]], dtype=float32)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X).sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0EyH21rUnwVx"
   },
   "source": [
    "### multiple regression\n",
    "Input dim is $D$, then there are $2D+2 = 2(D+1)$ parameters.\n",
    "\n",
    "$$\n",
    "p(y\\mid {\\bf x, w_\\mu, w_\\sigma}, b_\\mu, b_\\sigma)\n",
    "= \\mathcal N(y \\mid {\\bf w^T_\\mu x} + b_\\mu, {\\bf w^T_\\sigma x} + b_\\sigma)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oRMqNWPovp9x"
   },
   "outputs": [],
   "source": [
    "input_dim = 3\n",
    "output_dim = 1\n",
    "\n",
    "model = tfk.Sequential([\n",
    "    tfkl.InputLayer([input_dim]),\n",
    "    tfkl.Dense(tfpl.IndependentNormal.params_size(output_dim)),\n",
    "    tfpl.IndependentNormal(output_dim)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 233
    },
    "colab_type": "code",
    "id": "fJ57O24lvvFw",
    "outputId": "1892d53d-618c-432e-88dd-6e59065cda2c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_3\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense_2 (Dense)              (None, 2)                 8         \n",
      "_________________________________________________________________\n",
      "independent_normal_3 (Indepe ((None, 1), (None, 1))    0         \n",
      "=================================================================\n",
      "Total params: 8\n",
      "Trainable params: 8\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "f43-SJlIvw8v"
   },
   "outputs": [],
   "source": [
    "X = tf.constant([[3., -4., 1.],\n",
    "                 [2., 2., 2.],\n",
    "                 [-5., 3., 1.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "eiuaNMNkv4q6",
    "outputId": "74db0250-f38a-4a09-d5e1-5707d02096f0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.Independent 'sequential_3_independent_normal_3_IndependentNormal_Independentsequential_3_independent_normal_3_IndependentNormal_Normal' batch_shape=[3] event_shape=[1] dtype=float32>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 89
    },
    "colab_type": "code",
    "id": "LiP0lDahEQLU",
    "outputId": "57ce963d-b4b0-44c7-860e-4ca0651a2ab3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=2168, shape=(3, 1), dtype=float32, numpy=\n",
       "array([[-5.4155235],\n",
       "       [-2.0300615],\n",
       "       [ 3.971672 ]], dtype=float32)>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X).sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b8D3SXnMv6WW"
   },
   "source": [
    "### まとめ\n",
    "As you see, `tfp.layers` module is writen like `tf.keras.layers` module with `tf.keras.Sequential`.\n",
    "\n",
    "Note that, `tf.keras.layers` has `__call__` method whici recive `tf.Tensor` and output `tf.Tensor`, but `tfp.layers` output `tfp.distributions.Distribution` class. When you get some value type of `tf.Tensor`, you use `sample()`.\n",
    "\n",
    "If you want to `N` samples, you write `object.sample(N)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 377
    },
    "colab_type": "code",
    "id": "hoLAR6htF5qt",
    "outputId": "c7b0121e-ab7a-42c3-9c2f-cf40121b945b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=538, shape=(5, 3, 1), dtype=float32, numpy=\n",
       "array([[[ 1.6277615 ],\n",
       "        [-0.57028174],\n",
       "        [ 0.7915349 ]],\n",
       "\n",
       "       [[ 1.6302956 ],\n",
       "        [-2.5077786 ],\n",
       "        [ 1.7426147 ]],\n",
       "\n",
       "       [[ 1.5174006 ],\n",
       "        [-3.0163887 ],\n",
       "        [-0.4108941 ]],\n",
       "\n",
       "       [[ 1.7809215 ],\n",
       "        [-2.4185162 ],\n",
       "        [ 3.8518336 ]],\n",
       "\n",
       "       [[ 1.8309169 ],\n",
       "        [-4.1943836 ],\n",
       "        [-1.7214508 ]]], dtype=float32)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X).sample(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cxxm_bebHlrZ"
   },
   "source": [
    "This implementation is useful for bayes prediction distribution.\n",
    "\n",
    "When you get from below distribution,\n",
    "\n",
    "$$\n",
    "p(y\\mid x) = \\int _ w p(y \\mid w, x)p(w\\mid D)dw \n",
    "$$\n",
    "\n",
    "you can use below montecalro sampling.\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "w_i &\\sim p(w \\mid D)\\\\\n",
    "y_i &\\sim p(y\\mid w_i, x) \\\\\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "This calculation approximates the integral with the sum.\n",
    "\n",
    "$$\n",
    "p(y \\mid x) \\simeq \\frac{1}{N} \\sum _ {i=1} ^N p(y\\mid w_i, x) p(w_i \\mid D)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gVDTgLIqGApu"
   },
   "source": [
    "### Variational Bayes Multiple regression\n",
    "You design statistical model below, \n",
    "\n",
    "$$\n",
    "p(y\\mid {\\bf x, w_\\mu, w_\\sigma}, b_\\mu, b_\\sigma)\n",
    "= \\mathcal N(y \\mid {\\bf w^T_\\mu x} + b_\\mu, {\\bf w^T_\\sigma x} + b_\\sigma)\n",
    "$$\n",
    "\n",
    "and set prior below.\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "p({\\bf w_\\mu}) = \\mathcal N(0, {\\bf \\Sigma_{w_\\mu}})\\\\\n",
    "p({\\bf w_\\sigma}) = \\mathcal N(0, {\\bf \\Sigma_{w_\\sigma}})\\\\\n",
    "p(b_\\mu) = \\mathcal N(0, \\sigma_{b_\\mu})\\\\\n",
    "p(b_\\sigma) = \\mathcal N(0, \\sigma_{b_\\sigma})\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "In variational inference, you prepare variational distribution for approximate $q({\\bf w_\\mu}), q({\\bf w_\\sigma}), q(b_\\mu) , q(b_\\sigma)$, and when each distiributions are normal distribution, in case of input dim is $D$, there are $4D + 2$ parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "enyzCWC5NCTn"
   },
   "outputs": [],
   "source": [
    "input_dim = 3\n",
    "output_dim = 1\n",
    "\n",
    "model = tfk.Sequential([\n",
    "    tfkl.InputLayer([input_dim]),\n",
    "    tfpl.DenseReparameterization(tfpl.IndependentNormal.params_size(output_dim)),\n",
    "    tfpl.IndependentNormal(output_dim)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 233
    },
    "colab_type": "code",
    "id": "L4IgbfmBOm2b",
    "outputId": "f8999084-1c5a-4da9-b1f2-ae22c74d1bf7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_4\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense_reparameterization_1 ( (None, 2)                 14        \n",
      "_________________________________________________________________\n",
      "independent_normal_4 (Indepe ((None, 1), (None, 1))    0         \n",
      "=================================================================\n",
      "Total params: 14\n",
      "Trainable params: 14\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cTjznRQxOtRM"
   },
   "outputs": [],
   "source": [
    "X = tf.constant([[3., -4., 1.],\n",
    "                 [2., 2., 2.],\n",
    "                 [-5., 3., 1.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "M_FHfL81PNG_",
    "outputId": "f013d00f-9b94-41ca-8f87-4cce7f952d6f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.Independent 'sequential_4_independent_normal_4_IndependentNormal_Independentsequential_4_independent_normal_4_IndependentNormal_Normal' batch_shape=[3] event_shape=[1] dtype=float32>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "NYuYJMZZPN1O",
    "outputId": "be3a0d00-916b-4c09-e5c3-112b386e62cc"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=2694, shape=(3, 1), dtype=float32, numpy=\n",
       "array([[ 1.1869445 ],\n",
       "       [-0.29076326],\n",
       "       [-1.5732886 ]], dtype=float32)>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(X).sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "v_vuQsTePPaL"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "name": "probabilistic_model_by_tfp_layers.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}