{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jxv6goXm7oGF"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "llMNufAK7nfK"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8Byow2J6LaPl"
      },
      "source": [
        "# tf.function and AutoGraph in TensorFlow 2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kGXS3UWBBNoc"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/beta/guide/autograph\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/guide/autograph.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/autograph.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/site/en/r2/guide/autograph.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CydFK2CL7ZHA"
      },
      "source": [
        "TF 2.0 brings together the ease of eager execution and the power of TF 1.0. At the center of this merger is `tf.function`, which allows you to transform a subset of Python syntax into portable, high-performance TensorFlow graphs.\n",
        "\n",
        "A cool new feature of `tf.function` is AutoGraph, which lets you write graph code using natural Python syntax. For a list of the Python features that you can use with AutoGraph, see [AutoGraph Capabilities and Limitations](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/LIMITATIONS.md). For more details about `tf.function`, see the RFC [TF 2.0: Functions, not Sessions](https://github.com/tensorflow/community/blob/master/rfcs/20180918-functions-not-sessions-20.md). For more details about AutoGraph, see `tf.autograph`.\n",
        "\n",
        "This tutorial will walk you through the basic features of `tf.function` and AutoGraph."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n4EKOpw9mObL"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Import TensorFlow 2.0 Preview Nightly and enable TF 2.0 mode:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "V9oECvVSI1Kj"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mT7meGqrZTz9"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow==2.0.0-beta1\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "77AsVr1GGtBP"
      },
      "source": [
        "## The `tf.function` decorator\n",
        "\n",
        "When you annotate a function with `tf.function`, you can still call it like any other function. But it will be compiled into a graph, which means you get the benefits of faster execution, running on GPU or TPU, or exporting to SavedModel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FhIg7-z6HNWj"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def simple_nn_layer(x, y):\n",
        "  return tf.nn.relu(tf.matmul(x, y))\n",
        "\n",
        "\n",
        "x = tf.random.uniform((3, 3))\n",
        "y = tf.random.uniform((3, 3))\n",
        "\n",
        "simple_nn_layer(x, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U-LAE4pMNR9g"
      },
      "source": [
        "If we examine the result of the annotation, we can see that it's a special callable that handles all interactions with the TensorFlow runtime."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "q4t2iuS7Nqc0"
      },
      "outputs": [],
      "source": [
        "simple_nn_layer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DqeefLGNXjZQ"
      },
      "source": [
        "If your code uses multiple functions, you don't need to annotate them all - any functions called from an annotated function will also run in graph mode."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3VGF7tlVXiZY"
      },
      "outputs": [],
      "source": [
        "def linear_layer(x):\n",
        "  return 2 * x + 1\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def deep_net(x):\n",
        "  return tf.nn.relu(linear_layer(x))\n",
        "\n",
        "\n",
        "deep_net(tf.constant((1, 2, 3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yQvg6ZSKWyqE"
      },
      "source": [
        "Functions can be faster than eager code, for graphs with many small ops. But for graphs with a few expensive ops (like convolutions), you may not see much speedup.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0EL6lVwEWuFo"
      },
      "outputs": [],
      "source": [
        "import timeit\n",
        "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
        "\n",
        "@tf.function\n",
        "def conv_fn(image):\n",
        "  return conv_layer(image)\n",
        "\n",
        "image = tf.zeros([1, 200, 200, 100])\n",
        "# warm up\n",
        "conv_layer(image); conv_fn(image)\n",
        "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
        "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
        "print(\"Note how there's not much difference in performance for convolutions\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "L4zj-jpH0jKH"
      },
      "outputs": [],
      "source": [
        "lstm_cell = tf.keras.layers.LSTMCell(10)\n",
        "\n",
        "@tf.function\n",
        "def lstm_fn(input, state):\n",
        "  return lstm_cell(input, state)\n",
        "\n",
        "input = tf.zeros([10, 10])\n",
        "state = [tf.zeros([10, 10])] * 2\n",
        "# warm up\n",
        "lstm_cell(input, state); lstm_fn(input, state)\n",
        "print(\"eager lstm:\", timeit.timeit(lambda: lstm_cell(input, state), number=10))\n",
        "print(\"function lstm:\", timeit.timeit(lambda: lstm_fn(input, state), number=10))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ohbSnA79mcJV"
      },
      "source": [
        "## Use Python control flow\n",
        "\n",
        "When using data-dependent control flow inside `tf.function`, you can use Python control flow statements and AutoGraph will convert them into appropriate TensorFlow ops. For example, `if` statements will be converted into `tf.cond()` if they depend on a `Tensor`.\n",
        "\n",
        "In the example below, `x` is a `Tensor` but the `if` statement works as expected:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aA3gOodCBkOw"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square_if_positive(x):\n",
        "  if x \u003e 0:\n",
        "    x = x * x\n",
        "  else:\n",
        "    x = 0\n",
        "  return x\n",
        "\n",
        "\n",
        "print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))\n",
        "print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GMiCUkdyoq98"
      },
      "source": [
        "Note: The previous example uses simple conditionals with scalar values. \u003ca href=\"#batching\"\u003eBatching\u003c/a\u003e is typically used in real-world code."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m-jWmsCmByyw"
      },
      "source": [
        "AutoGraph supports common Python statements like `while`, `for`, `if`, `break`, `continue` and `return`, with support for nesting. That means you can use `Tensor` expressions in the condition of `while` and `if` statements, or iterate over a `Tensor` in a `for` loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "toxKBOXbB1ro"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def sum_even(items):\n",
        "  s = 0\n",
        "  for c in items:\n",
        "    if c % 2 \u003e 0:\n",
        "      continue\n",
        "    s += c\n",
        "  return s\n",
        "\n",
        "\n",
        "sum_even(tf.constant([10, 12, 15, 20]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AtDaLrbySw4j"
      },
      "source": [
        "AutoGraph also provides a low-level API for advanced users. For example we can use it to have a look at the generated code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aRsde3x_SjTQ"
      },
      "outputs": [],
      "source": [
        "print(tf.autograph.to_code(sum_even.python_function))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rvJXCfk8VkLf"
      },
      "source": [
        "Here's an example of more complicated control flow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "h-Z87IJqVlKl"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  msg = tf.constant('')\n",
        "  for i in tf.range(n):\n",
        "    if tf.equal(i % 3, 0):\n",
        "      tf.print('Fizz')\n",
        "    elif tf.equal(i % 5, 0):\n",
        "      tf.print('Buzz')\n",
        "    else:\n",
        "      tf.print(i)\n",
        "\n",
        "fizzbuzz(tf.constant(15))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "h_Y4uC1R1B55"
      },
      "source": [
        "## Keras and AutoGraph\n",
        "\n",
        "AutoGraph is available by default in non-dynamic Keras models. For more information, see `tf.keras`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cR6mpLKP1HLe"
      },
      "outputs": [],
      "source": [
        "class CustomModel(tf.keras.models.Model):\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, input_data):\n",
        "    if tf.reduce_mean(input_data) \u003e 0:\n",
        "      return input_data\n",
        "    else:\n",
        "      return input_data // 2\n",
        "\n",
        "\n",
        "model = CustomModel()\n",
        "\n",
        "model(tf.constant([-2, -4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NTEvpBK9f8kj"
      },
      "source": [
        "## Side effects\n",
        "\n",
        "Just like in eager mode, you can use operations with side effects, like `tf.assign` or `tf.print` normally inside `tf.function`, and it will insert the necessary control dependencies to ensure they execute in order."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-Wd6i8S9gcuC"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(5)\n",
        "\n",
        "@tf.function\n",
        "def find_next_odd():\n",
        "  v.assign(v + 1)\n",
        "  if tf.equal(v % 2, 0):\n",
        "    v.assign(v + 1)\n",
        "\n",
        "\n",
        "find_next_odd()\n",
        "v"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4LfnJjm0Bm0B"
      },
      "source": [
        "## Example: training a simple model\n",
        "\n",
        "AutoGraph also allows you to move more computation inside TensorFlow. For example, a training loop is just control flow, so it can actually be brought into TensorFlow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Em5dzSUOtLRP"
      },
      "source": [
        "### Download data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "xqoxumv0ssQW"
      },
      "outputs": [],
      "source": [
        "def prepare_mnist_features_and_labels(x, y):\n",
        "  x = tf.cast(x, tf.float32) / 255.0\n",
        "  y = tf.cast(y, tf.int64)\n",
        "  return x, y\n",
        "\n",
        "def mnist_dataset():\n",
        "  (x, y), _ = tf.keras.datasets.mnist.load_data()\n",
        "  ds = tf.data.Dataset.from_tensor_slices((x, y))\n",
        "  ds = ds.map(prepare_mnist_features_and_labels)\n",
        "  ds = ds.take(20000).shuffle(20000).batch(100)\n",
        "  return ds\n",
        "\n",
        "train_dataset = mnist_dataset()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "znmy4l8ntMvW"
      },
      "source": [
        "### Define the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ltxyJVWTqNAO"
      },
      "outputs": [],
      "source": [
        "model = tf.keras.Sequential((\n",
        "    tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(10)))\n",
        "model.build()\n",
        "optimizer = tf.keras.optimizers.Adam()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oeYV6mKnJGMr"
      },
      "source": [
        "### Define the training loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3xtg_MMhJETd"
      },
      "outputs": [],
      "source": [
        "compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "\n",
        "compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "\n",
        "def train_one_step(model, optimizer, x, y):\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = model(x)\n",
        "    loss = compute_loss(y, logits)\n",
        "\n",
        "  grads = tape.gradient(loss, model.trainable_variables)\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "\n",
        "  compute_accuracy(y, logits)\n",
        "  return loss\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def train(model, optimizer):\n",
        "  train_ds = mnist_dataset()\n",
        "  step = 0\n",
        "  loss = 0.0\n",
        "  accuracy = 0.0\n",
        "  for x, y in train_ds:\n",
        "    step += 1\n",
        "    loss = train_one_step(model, optimizer, x, y)\n",
        "    if tf.equal(step % 10, 0):\n",
        "      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())\n",
        "  return step, loss, accuracy\n",
        "\n",
        "step, loss, accuracy = train(model, optimizer)\n",
        "print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SnsumiP6eRYL"
      },
      "source": [
        "## Batching\n",
        "\n",
        "In real applications batching is essential for performance. The best code to convert to AutoGraph is code where the control flow is decided at the _batch_ level. If making decisions at the individual _example_ level, try to use batch APIs to maintain performance.\n",
        "\n",
        "For example, if you have the following code in Python:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t31QoERiNccJ"
      },
      "outputs": [],
      "source": [
        "def square_if_positive(x):\n",
        "  return [i ** 2 if i \u003e 0 else i for i in x]\n",
        "\n",
        "\n",
        "square_if_positive(range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kSeEJ76uNgwD"
      },
      "source": [
        "You may be tempted to write it in TensorFlow as such (and this would work!):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RqR8WzSzNf87"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square_if_positive_naive(x):\n",
        "  result = tf.TensorArray(tf.int32, size=x.shape[0])\n",
        "  for i in tf.range(x.shape[0]):\n",
        "    if x[i] \u003e 0:\n",
        "      result = result.write(i, x[i] ** 2)\n",
        "    else:\n",
        "      result = result.write(i, x[i])\n",
        "  return result.stack()\n",
        "\n",
        "\n",
        "square_if_positive_naive(tf.range(-5, 5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gTcyWXVGN3gS"
      },
      "source": [
        "But in this case, it turns out you can write the following:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VO2f6x-lNfVj"
      },
      "outputs": [],
      "source": [
        "def square_if_positive_vectorized(x):\n",
        "  return tf.where(x \u003e 0, x ** 2, x)\n",
        "\n",
        "\n",
        "square_if_positive_vectorized(tf.range(-5, 5))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Jxv6goXm7oGF"
      ],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "TF2.0 Functions and Autograph",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Hl4PR32Y_nhq4AZI-ku9FlpZIgXjEU7d",
          "timestamp": 1544141357531
        },
        {
          "file_id": "https://github.com/tensorflow/docs/blob/master/site/en/guide/autograph.ipynb",
          "timestamp": 1542126584672
        }
      ],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}