TensorFlow: Advanced Techniques/C2_W3_Lab_2-graphs-for-complex-code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/https-deeplearning-ai/tensorflow-3-public/blob/main/Course%202%20-%20Custom%20Training%20loops%2C%20Gradients%20and%20Distributed%20Training/Week%203%20-%20Autograph/C2_W3_Lab_2-graphs-for-complex-code.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hrXv0rU9sIma"
   },
   "source": [
    "# Autograph: Graphs for complex code\n",
    "\n",
    "In this ungraded lab, you'll go through some of the scenarios from the lesson `Creating graphs for complex code`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NiolgWMPgpwI"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    # %tensorflow_version only exists in Colab.\n",
    "    %tensorflow_version 2.x\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you saw in the lectures, seemingly simple functions can sometimes be difficult to write in graph mode. Fortunately, Autograph generates this complex graph code for us.\n",
    "\n",
    "- Here is a function that does some multiplication and additon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = tf.Variable(1.0)\n",
    "b = tf.Variable(2.0)\n",
    "\n",
    "@tf.function\n",
    "def f(x,y):\n",
    "    a.assign(y * b)\n",
    "    b.assign_add(x * a)\n",
    "    return a + b\n",
    "\n",
    "print(f(1.0, 2.0))\n",
    "\n",
    "print(tf.autograph.to_code(f.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Here is a function that checks if the sign of a number is positive or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def sign(x):\n",
    "    if x > 0:\n",
    "        return 'Positive'\n",
    "    else:\n",
    "        return 'Negative or zero'\n",
    "\n",
    "print(\"Sign = {}\".format(sign(tf.constant(2))))\n",
    "print(\"Sign = {}\".format(sign(tf.constant(-2))))\n",
    "\n",
    "print(tf.autograph.to_code(sign.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Here is another function that includes a while loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    while tf.reduce_sum(x) > 1:\n",
    "        tf.print(x)\n",
    "        x = tf.tanh(x)\n",
    "    return x\n",
    "\n",
    "print(tf.autograph.to_code(f.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Here is a function that uses a for loop and an if statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def sum_even(items):\n",
    "    s = 0\n",
    "    for c in items:\n",
    "        if c % 2 > 0:\n",
    "            continue\n",
    "        s += c\n",
    "    return s\n",
    "\n",
    "print(tf.autograph.to_code(sum_even.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Print statements\n",
    "\n",
    "Tracing also behaves differently in graph mode. First, here is a function (not decorated with `@tf.function` yet) that prints the value of the input parameter.  `f(2)` is called in a for loop 5 times, and then `f(3)` is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    print(\"Traced with\", x)\n",
    "\n",
    "for i in range(5):\n",
    "    f(2)\n",
    "    \n",
    "f(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you were to decorate this function with `@tf.function` and run it, notice that the print statement only appears once for `f(2)` even though it is called in a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    print(\"Traced with\", x)\n",
    "\n",
    "for i in range(5):\n",
    "    f(2)\n",
    "    \n",
    "f(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now compare `print` to `tf.print`.\n",
    "- `tf.print` is graph aware and will run as expected in loops. \n",
    "\n",
    "Try running the same code where `tf.print()` is added in addition to the regular `print`.\n",
    "- Note how `tf.print` behaves compared to `print` in graph mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    print(\"Traced with\", x)\n",
    "    # added tf.print\n",
    "    tf.print(\"Executed with\", x)\n",
    "\n",
    "for i in range(5):\n",
    "    f(2)\n",
    "    \n",
    "f(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Avoid defining variables inside the function\n",
    "\n",
    "This function (not decorated yet) defines a tensor `v` and adds the input `x` to it.  \n",
    "\n",
    "Here, it runs fine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    v = tf.Variable(1.0)\n",
    "    v.assign_add(x)\n",
    "    return v\n",
    "\n",
    "print(f(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if you decorate the function with `@tf.function`.\n",
    "\n",
    "The cell below will throw an error because `tf.Variable` is defined within the function. The graph mode function should only contain operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    v = tf.Variable(1.0)\n",
    "    v.assign_add(x)\n",
    "    return v\n",
    "\n",
    "print(f(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get around the error above, simply move `v = tf.Variable(1.0)` to the top of the cell before the `@tf.function` decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define the variables outside of the decorated function\n",
    "v = tf.Variable(1.0)\n",
    "\n",
    "@tf.function\n",
    "def f(x):\n",
    "    return v.assign_add(x)\n",
    "\n",
    "print(f(5))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Autograph Lesson 2.ipynb",
   "private_outputs": true,
   "provenance": []
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}