Tensorflow Differentiation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Taking derivatives with Tensorflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this short notebook, we'll look at some examples of how to do unusual derivatives with tensorflow.  Some of this is useful in the \"real\"  notebooks, and some of it is irrelevant to the AI4NP school but may be useful to you elsewhere."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example from the tensorflow Gradient Tape example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(6.0, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "x = tf.constant(3.0)\n",
    "with tf.GradientTape() as g:\n",
    "  g.watch(x)\n",
    "  y = x * x\n",
    "dy_dx = g.gradient(y, x)\n",
    "print(dy_dx)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: We had to make sure to mark that we wanted a derivative with respect to x!  That is because tensorflow only cares (by default) about tf.Variable objects.\n",
    "\n",
    "This is because NN weights are typically tf.Variables, and most users don't want the gradient with respect to the input.  In our applications, though, we absolutely want that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "x = tf.constant(3.0)\n",
    "with tf.GradientTape() as g:\n",
    "  y = x * x\n",
    "dy_dx = g.gradient(y, x)\n",
    "print(dy_dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without declaring it, we get `None`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(6.0, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "x = tf.Variable(3.0, trainable=True)\n",
    "with tf.GradientTape() as g:\n",
    "  y = x * x\n",
    "dy_dx = g.gradient(y, x)\n",
    "print(dy_dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If x is a variable, we're all good again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiple Tapes\n",
    "\n",
    "This example is from (Advanced Automatic Differentiation)[https://www.tensorflow.org/guide/advanced_autodiff]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = tf.constant(0.0)\n",
    "x1 = tf.constant(0.0)\n",
    "\n",
    "with tf.GradientTape() as tape0, tf.GradientTape() as tape1:\n",
    "  tape0.watch(x0)\n",
    "  tape1.watch(x1)\n",
    "\n",
    "  y0 = tf.math.sin(x0)\n",
    "  y1 = tf.nn.sigmoid(x1)\n",
    "\n",
    "  y = y0 + y1\n",
    "\n",
    "  ys = tf.reduce_sum(y)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tape0.gradient(ys, x0).numpy()   # cos(x) => 1.0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.25"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TF Functions\n",
    "\n",
    "Are tf functions really faster?  Let's find out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poorly_written_function(x, y):\n",
    "    x = x + 3\n",
    "    x = tf.exp(x)\n",
    "    x = x - 7\n",
    "    y = y**2\n",
    "    y = tf.math.tanh(y)\n",
    "    y = y -1\n",
    "    \n",
    "    return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.convert_to_tensor(numpy.arange(0, 100., 0.01))\n",
    "y = tf.convert_to_tensor(numpy.arange(-50.,50, 0.01))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = poorly_written_function(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.005386841017752886\n"
     ]
    }
   ],
   "source": [
    "import timeit\n",
    "\n",
    "t = timeit.Timer(lambda: poorly_written_function(x,y))  \n",
    "print(t.timeit(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def better_written_function(x,y):\n",
    "    x = tf.exp(x + 3.) - 7.\n",
    "    y = tf.math.tanh(y**2) - 1\n",
    "    y = y -1\n",
    "        \n",
    "    return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.004837962798774242\n"
     ]
    }
   ],
   "source": [
    "import timeit\n",
    "\n",
    "t = timeit.Timer(lambda: better_written_function(x,y))  \n",
    "print(t.timeit(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Things went a little faster here because we didn't go into python as much."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10000,), dtype=float64, numpy=\n",
       "array([1.30855369e+01, 1.32873999e+01, 1.34912917e+01, ...,\n",
       "       5.23965632e+44, 5.29231574e+44, 5.34550440e+44])>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "traced_function = tf.function(poorly_written_function)\n",
    "traced_function(x,y) # the very first call does the trace and graph compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0027769929729402065\n"
     ]
    }
   ],
   "source": [
    "import timeit\n",
    "\n",
    "t = timeit.Timer(lambda: traced_function(x,y))  \n",
    "print(t.timeit(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even faster this time!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even for this (useless) function, we can compute derivatives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    tape.watch(x)\n",
    "    tape.watch(y)\n",
    "    z = traced_function(x,y)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[2.00855369e+01 2.02873999e+01 2.04912917e+01 ... 5.23965632e+44\n",
      " 5.29231574e+44 5.34550440e+44], shape=(10000,), dtype=float64)\n"
     ]
    }
   ],
   "source": [
    "dzdx = tape.gradient(z, x)\n",
    "print(dzdx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([0. 0. 0. ... 0. 0. 0.], shape=(10000,), dtype=float64)\n"
     ]
    }
   ],
   "source": [
    "dzdy = tape.gradient(z, y)\n",
    "print(dzdy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Is it right?  Well:\n",
    "$z = tanh(y**2) - 1 + e^{x + 3} -7$\n",
    "\n",
    "$\\frac{dz}{dx} = \\frac{d}{dx} e^{(x+3)} = e^3 * \\frac{d}{dx} e^x $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10000,), dtype=float64, numpy=\n",
       "array([0.00000000e+00, 3.55271368e-15, 0.00000000e+00, ...,\n",
       "       0.00000000e+00, 0.00000000e+00, 0.00000000e+00])>"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "numpy.exp(3.) * tf.exp(x) - dzdx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks pretty good!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}