02b_nll_tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting: Computing an NLL\n",
    "\n",
    "We will be using Tensorflow's new eager mode, the new JIT static graph, and a classic API static graph to solve a different problem: fitting unbinned datasets. Like before, let's set up the data and then try a solution with NumPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "dist = np.hstack(\n",
    "    [\n",
    "        np.random.normal(loc=1, scale=2.0, size=500_000),\n",
    "        np.random.normal(loc=1, scale=0.5, size=500_000),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's load TensorFlow 2.0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "print(f\"{tf.__version__ = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset does not change, so that can be a constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_dist = tf.constant(dist, name=\"dist\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that this looks a lot like NumPy, except most of the names are different. Also this is the same on both APIs; the main difference is the setup and debugging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian(x, μ, σ):\n",
    "    return 1 / tf.sqrt(2 * np.pi * σ**2) * tf.math.exp(-((x - μ) ** 2) / (2 * σ**2))\n",
    "\n",
    "\n",
    "def add(x, f_0, mean, sigma, sigma2):\n",
    "    return f_0 * gaussian(x, mean, sigma) + (1 - f_0) * gaussian(x, mean, sigma2)\n",
    "\n",
    "\n",
    "def make_nll(dist, f_0, mean, sigma, sigma2):\n",
    "    return -tf.reduce_sum(tf.math.log(add(dist, f_0, mean, sigma, sigma2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "make_nll(d_dist, *np.random.rand(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try using the autograph technique to convert this into something like a static graph (it gets cached on first use). This could be written as a decorator, `@tf.function`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nll = tf.function(make_nll)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the static graph to work, we need to be careful and use all TensorFlow objects:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_f_0 = tf.Variable(0, dtype=tf.float64)\n",
    "tf_mean = tf.Variable(0, dtype=tf.float64)\n",
    "tf_sigma = tf.Variable(0, dtype=tf.float64)\n",
    "tf_sigma2 = tf.Variable(0, dtype=tf.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "r = np.random.rand(4)\n",
    "\n",
    "tf_f_0.assign(r[0])\n",
    "tf_mean.assign(r[1])\n",
    "tf_sigma.assign(r[1])\n",
    "tf_sigma2.assign(r[1])\n",
    "\n",
    "nll(d_dist, tf_f_0, tf_mean, tf_sigma, tf_sigma2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Static Graph (classic API)\n",
    "\n",
    "Let's try the classic API, and build a static graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow.compat.v1 as tf\n",
    "\n",
    "tf.disable_eager_execution()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Repeating this here for good measure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian(x, μ, σ):\n",
    "    return 1 / tf.sqrt(2 * np.pi * σ**2) * tf.math.exp(-((x - μ) ** 2) / (2 * σ**2))\n",
    "\n",
    "\n",
    "def add(x, f_0, mean, sigma, sigma2):\n",
    "    return f_0 * gaussian(x, mean, sigma) + (1 - f_0) * gaussian(x, mean, sigma2)\n",
    "\n",
    "\n",
    "def make_nll(dist, f_0, mean, sigma, sigma2):\n",
    "    return -tf.reduce_sum(tf.math.log(add(dist, f_0, mean, sigma, sigma2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The API is different (we still have constant, but placeholder is classic API):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.constant(dist)\n",
    "tf_f_0 = tf.placeholder(dtype=tf.float64)\n",
    "tf_mean = tf.placeholder(dtype=tf.float64)\n",
    "tf_sigma = tf.placeholder(dtype=tf.float64)\n",
    "tf_sigma2 = tf.placeholder(dtype=tf.float64)\n",
    "\n",
    "graph = make_nll(x, tf_f_0, tf_mean, tf_sigma, tf_sigma2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to start up a session, then feed the hungry graph with values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = tf.Session()\n",
    "\n",
    "\n",
    "def nll(f_0, mean, sigma, sigma2):\n",
    "    loss_value = sess.run(\n",
    "        graph,\n",
    "        feed_dict={tf_f_0: f_0, tf_mean: mean, tf_sigma: sigma, tf_sigma2: sigma2},\n",
    "    )\n",
    "    return loss_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "nll(*np.random.rand(4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyGPU Course 2019/12 [course/pygpu/default]",
   "language": "python",
   "name": "sys_pygpu201912"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}