Tutorial6.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AutoGraph\n",
    "\n",
    "In this tutorial, we will introduce how to convert python function into executable\n",
    "Tensorflow graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "import time\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.4.1\n"
     ]
    }
   ],
   "source": [
    "print(tf.__version__)\n",
    "tf.random.set_seed(42)\n",
    "np.random.seed(42)\n",
    "true_weights = tf.constant(list(range(5)), dtype=tf.float32)[:, tf.newaxis]\n",
    "x = tf.constant(tf.random.uniform((32, 5)), dtype=tf.float32)\n",
    "y = tf.constant(x @ true_weights, dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AutoGraph\n",
    "\n",
    "A computational graph contains two things: computation and data. Tensorflow graph `tf.Graph` is the computational graph made from operations as computation units and tensors as data units. Tensorflow has many optimizations around graphs, so executing in graph mode result in better utilization of distributed computing. Instead of writing complicated graph mode code, Tensorflow 2 provides a tool AutoGraph to automatically analyze and convert python code into graph code which can then be traced to create a graph with Function.\n",
    "\n",
    "The following python function involves relatively simple math operations on tensors, so it is pretty much graph ready. Let's see if `tf.autograph` would do anything interesting to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        def tf__f(a, b, power=None, d=None):\n",
      "            with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "                do_return = False\n",
      "                retval_ = ag__.UndefinedReturnValue()\n",
      "                try:\n",
      "                    do_return = True\n",
      "                    retval_ = (ag__.converted_call(ag__.ld(tf).pow, (ag__.ld(a), ag__.ld(power)), None, fscope) + (ag__.ld(d) * ag__.ld(b)))\n",
      "                except:\n",
      "                    do_return = False\n",
      "                    raise\n",
      "                return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def f(a, b, power=2, d=3):\n",
    "    return tf.pow(a, power) + d * b\n",
    "\n",
    "converted_f = tf.autograph.to_graph(f)\n",
    "print(inspect.getsource(converted_f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def f(a, b, power=2, d=3):\n",
      "    return tf.pow(a, power) + d * b\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(f))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's move on to another python function with a bit graph unfriendly construct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        def tf__cube(x):\n",
      "            with ag__.FunctionScope('cube', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "                do_return = False\n",
      "                retval_ = ag__.UndefinedReturnValue()\n",
      "                o = ag__.ld(x)\n",
      "\n",
      "                def get_state():\n",
      "                    return (o,)\n",
      "\n",
      "                def set_state(vars_):\n",
      "                    nonlocal o\n",
      "                    (o,) = vars_\n",
      "\n",
      "                def loop_body(itr):\n",
      "                    nonlocal o\n",
      "                    _ = itr\n",
      "                    o = ag__.ld(o)\n",
      "                    o *= x\n",
      "                _ = ag__.Undefined('_')\n",
      "                ag__.for_stmt(ag__.converted_call(ag__.ld(range), (2,), None, fscope), None, loop_body, get_state, set_state, ('o',), {'iterate_names': '_'})\n",
      "                try:\n",
      "                    do_return = True\n",
      "                    retval_ = ag__.ld(o)\n",
      "                except:\n",
      "                    do_return = False\n",
      "                    raise\n",
      "                return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def cube(x):\n",
    "    o = x\n",
    "    for _ in range(2):\n",
    "        o *= x\n",
    "    return o\n",
    "\n",
    "converted_cube = tf.autograph.to_graph(cube)\n",
    "print(inspect.getsource(converted_cube))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that body of the for loop is converted into a *loop_body* function, which is invoked by *autograph.for_stmt*. This `for_stmt` operation is, in a sense, overloads the for statement. The purpose of this transformation is to make the code into a functional style so that it can be executed in graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        def tf__g(x):\n",
      "            with ag__.FunctionScope('g', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "                do_return = False\n",
      "                retval_ = ag__.UndefinedReturnValue()\n",
      "\n",
      "                def get_state():\n",
      "                    return (do_return, retval_)\n",
      "\n",
      "                def set_state(vars_):\n",
      "                    nonlocal do_return, retval_\n",
      "                    (do_return, retval_) = vars_\n",
      "\n",
      "                def if_body():\n",
      "                    nonlocal do_return, retval_\n",
      "                    try:\n",
      "                        do_return = True\n",
      "                        retval_ = ag__.converted_call(ag__.ld(tf).square, (ag__.ld(x),), None, fscope)\n",
      "                    except:\n",
      "                        do_return = False\n",
      "                        raise\n",
      "\n",
      "                def else_body():\n",
      "                    nonlocal do_return, retval_\n",
      "                    try:\n",
      "                        do_return = True\n",
      "                        retval_ = ag__.ld(x)\n",
      "                    except:\n",
      "                        do_return = False\n",
      "                        raise\n",
      "                ag__.if_stmt(ag__.converted_call(ag__.ld(tf).reduce_any, ((ag__.ld(x) < 0),), None, fscope), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)\n",
      "                return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def g(x):\n",
    "    if tf.reduce_any(x < 0):\n",
    "        return tf.square(x)\n",
    "    return x\n",
    "\n",
    "converted_g = tf.autograph.to_graph(g)\n",
    "print(inspect.getsource(converted_g))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The big idea about AutoGraph is that it translates the python code we wrote into a style that can be traced to create Tensorflow graphs. During the transformation, great efforts were made to rewrite the data-dependent conditionals and control flows, as these statements can't be straight forward overloaded by Tensorflow operations. There are a lot more about AutoGraph, the interested reader can refer to the [paper](https://arxiv.org/abs/1810.08061) for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions\n",
    "\n",
    "Once the code is graph friendly, we can trace the operations in the code to create a graph. The created graph then gets wrapped up in a `ConcreteFunction` object so that we can execute computations backed in graph mode with it. \n",
    "\n",
    "Lets walkthrough how these works. First we provide Tensorflow our graph friendly code to\n",
    "create a `Function` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "tf_func_f = tf.function(autograph=False)(f)\n",
    "tf_func_g = tf.function(autograph=False)(converted_g)\n",
    "tf_func_g2 = tf.function(autograph=True)(g)\n",
    "print(tf_func_f.python_function is f)\n",
    "print(tf_func_g.python_function is converted_g)\n",
    "print(tf_func_g2.python_function is g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As of now, there is no graph created yet. We only attached our function to this Functionobject. Notice that we specified `autograph=False` when constructing the Function in the first two cases, because we know both f and converted_g are graph friendly. But good old g is not, so we need to turn on autograph for it. In fact,\n",
    "`tf.function(autograph=False)(tf.autograph.to_graph(g))` is roughly equivlent to `tf.function(autograph=True)(g)`.\n",
    "\n",
    "Note that we can create Function with `tf.function(autograph=False)(g)` this will\n",
    "succeed without error. But we won't be able to create any graphs with this Function in the\n",
    "next step.\n",
    "\n",
    "The next step is to provide Tensorflow a signature, i.e. a description of inputs, allowing it\n",
    "to create a graph by tracing how the input tensors would flow through the operations in\n",
    "the graph and record the graph in a callable object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ConcreteFunction tf__g(x)\n",
      "  Args:\n",
      "    x: float32 Tensor, shape=(3,)\n",
      "  Returns:\n",
      "    float32 Tensor, shape=(3,)\n"
     ]
    }
   ],
   "source": [
    "concrete_g = tf_func_g.get_concrete_function(x=tf.TensorSpec(shape=[3], dtype=tf.float32))\n",
    "print(concrete_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this concrete function directly as if it is a operation shipped with Tensorflow,\n",
    "or we can call the Function object, which will look up the concrete function and use it.\n",
    "Either way, the computation will be executed with the created graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<tf.Tensor: shape=(3,), dtype=float32, numpy=array([1., 1., 4.], dtype=float32)>\n",
      "<tf.Tensor: shape=(3,), dtype=float32, numpy=array([1., 1., 4.], dtype=float32)>\n"
     ]
    }
   ],
   "source": [
    "pprint(concrete_g(tf.constant([-1, 1, -2], dtype=tf.float32)))\n",
    "pprint(tf_func_g(tf.constant([-1, 1, -2], dtype=tf.float32)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Function object is like a graph factory. When detailed input specifications were\n",
    "provided, it uses the graph code as receipt to create new graphs. When asked with an\n",
    "known specifications, it will dig up the graph in the storage and serve it. When called with\n",
    "an unknown signature, it will trigger the creation of the concrete function first. Let's try to\n",
    "make a bunch graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ConcreteFunction f(a, b, power=2, d=3)\n",
      "  Args:\n",
      "    a: float32 Tensor, shape=(1,)\n",
      "    b: float32 Tensor, shape=(1,)\n",
      "  Returns:\n",
      "    float32 Tensor, shape=(1,)\n",
      "<tf.Tensor: shape=(), dtype=float32, numpy=7.0>\n",
      "<tf.Tensor: shape=(), dtype=float32, numpy=7.0>\n",
      "<tf.Tensor: shape=(), dtype=float32, numpy=7.0>\n",
      "<tf.Tensor: shape=(), dtype=float32, numpy=7.0>\n",
      "<tf.Tensor: shape=(), dtype=float32, numpy=7.0>\n"
     ]
    }
   ],
   "source": [
    "concrete_f = tf_func_f.get_concrete_function(a=tf.TensorSpec(shape=[1], dtype=tf.float32), b=tf.TensorSpec(shape=[1], dtype=tf.float32))\n",
    "print(concrete_f)\n",
    "pprint(concrete_f(tf.constant(1.), tf.constant(2.)))\n",
    "pprint(tf_func_f(1., 2.))\n",
    "pprint(tf_func_f(a=tf.constant(1., dtype=tf.float32), b=2, power=2.))\n",
    "pprint(tf_func_f(a=tf.constant(1., dtype=tf.float32), b=2., d=3))\n",
    "pprint(tf_func_f(a=tf.constant(1., dtype=tf.float32), b=2., d=3., power=3.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    }
   ],
   "source": [
    "## Print out the number of graphs we have created above\n",
    "print(tf_func_f._get_tracing_count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 ((TensorSpec(shape=(1,), dtype=tf.float32, name='a'), TensorSpec(shape=(1,), dtype=tf.float32, name='b'), 2, 3), {})\n",
      "1 ((1.0, 2.0, 2, 3), {})\n",
      "2 ((TensorSpec(shape=(), dtype=tf.float32, name='a'), 2, 2.0, 3), {})\n",
      "3 ((TensorSpec(shape=(), dtype=tf.float32, name='a'), 2.0, 3.0, 3.0), {})\n"
     ]
    }
   ],
   "source": [
    "### ??\n",
    "for i, f in enumerate(tf_func_f._list_all_concrete_functions_for_serialization()):\n",
    "    print(i, f.structured_input_signature)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, `tf.function` is also available as a decorator, which makes life easier. The\n",
    "following two are equivalent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function(autograph=False)\n",
    "def square(x):\n",
    "    return x * x\n",
    "\n",
    "square = tf.function(autograph=False)(square)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@tf.function(autograph=False)\n",
      "def square(x):\n",
      "    return x * x\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(square))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear Regression Revisited\n",
    "\n",
    "**Normal mode**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean squared loss at iteration    0 is 9.5217\n",
      "mean squared loss at iteration  200 is 0.0708\n",
      "mean squared loss at iteration  400 is 0.0082\n",
      "mean squared loss at iteration  600 is 0.0011\n",
      "mean squared loss at iteration  800 is 0.0002\n",
      "mean squared loss at iteration 1000 is 0.0000\n",
      "<tf.Variable 'Variable:0' shape=(5, 1) dtype=float32, numpy=\n",
      "array([[0.01446281],\n",
      "       [1.006089  ],\n",
      "       [1.9856374 ],\n",
      "       [2.9991345 ],\n",
      "       [3.992287  ]], dtype=float32)>\n",
      "time took: 0.6380012035369873 seconds\n"
     ]
    }
   ],
   "source": [
    "# ground truth\n",
    "true_weights = tf.constant(list(range(5)), dtype=tf.float32)[:, tf.newaxis]\n",
    "\n",
    "# some random training data\n",
    "x = tf.constant(tf.random.uniform((32, 5)), dtype=tf.float32)\n",
    "y = tf.constant(x @ true_weights, dtype=tf.float32)\n",
    "\n",
    "t0 = time.time()\n",
    "\n",
    "weights = tf.Variable(tf.random.uniform((5, 1)), dtype=tf.float32)\n",
    "for iteration in range(1001):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_hat = tf.linalg.matmul(x, weights)\n",
    "        loss = tf.reduce_mean(tf.square(y - y_hat))\n",
    "    if not (iteration % 200):\n",
    "        print('mean squared loss at iteration {:4d} is {:5.4f}'.format(iteration, loss))\n",
    "    gradients = tape.gradient(loss, weights)\n",
    "    weights.assign_add(-0.05 * gradients)\n",
    "\n",
    "pprint(weights)\n",
    "print('time took: {} seconds'.format(time.time() - t0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**In graph mode**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean squared loss at iteration    0 is 10.4993\n",
      "mean squared loss at iteration  200 is 0.0646\n",
      "mean squared loss at iteration  400 is 0.0079\n",
      "mean squared loss at iteration  600 is 0.0011\n",
      "mean squared loss at iteration  800 is 0.0002\n",
      "mean squared loss at iteration 1000 is 0.0000\n",
      "<tf.Variable 'Variable:0' shape=(5, 1) dtype=float32, numpy=\n",
      "array([[0.01437845],\n",
      "       [1.0059469 ],\n",
      "       [1.9856571 ],\n",
      "       [2.999293  ],\n",
      "       [3.992368  ]], dtype=float32)>\n",
      "time took: 0.27387380599975586 seconds\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "weights = tf.Variable(tf.random.uniform((5, 1)), dtype=tf.float32)\n",
    "@tf.function\n",
    "def train_step():\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_hat = tf.linalg.matmul(x, weights)\n",
    "        loss = tf.reduce_mean(tf.square(y - y_hat))\n",
    "    gradients = tape.gradient(loss, weights)\n",
    "    weights.assign_add(-0.05 * gradients)\n",
    "    return loss\n",
    "\n",
    "for iteration in range(1001):\n",
    "    loss = train_step()\n",
    "    if not (iteration % 200):\n",
    "        print('mean squared loss at iteration {:4d} is {:5.4f}'.format(iteration, loss))\n",
    "        \n",
    "pprint(weights)\n",
    "print('time took: {} seconds'.format(time.time() - t0))"
   ]
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}