tf_function_retracing_error.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tf_function_retracing_error.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyNGtQ8xzZMFIfIFagWmqVK4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jvishnuvardhan/Stackoverflow_Questions/blob/master/tf_function_retracing_error.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2TSDkr58XkUd",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 632
        },
        "outputId": "e2bb3e99-93d6-4016-a7cb-061152a15d8c"
      },
      "source": [
        "!pip install tf-nightly"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: tf-nightly in /usr/local/lib/python3.6/dist-packages (2.2.0.dev20200506)\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.6.3)\n",
            "Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.18.3)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.9.0)\n",
            "Requirement already satisfied: h5py<2.11.0,>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (2.10.0)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.12.1)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.12.0)\n",
            "Requirement already satisfied: scipy==1.4.1; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.4.1)\n",
            "Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.3.3)\n",
            "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.34.2)\n",
            "Requirement already satisfied: tf-estimator-nightly in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (2.3.0.dev2020050601)\n",
            "Requirement already satisfied: tb-nightly<2.4.0a0,>=2.3.0a0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (2.3.0a20200506)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (3.2.1)\n",
            "Requirement already satisfied: keras-preprocessing>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.1.0)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.28.1)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (3.10.0)\n",
            "Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.2.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.1.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.2.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (0.4.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.6.0.post3)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.0.1)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (2.23.0)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (46.1.3)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.7.2)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.3.0)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (2020.4.5.1)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (2.9)\n",
            "Requirement already satisfied: cachetools<3.2,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.1.1)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (0.2.8)\n",
            "Requirement already satisfied: rsa<4.1,>=3.1.4 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (4.0)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.1.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.6/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (0.4.8)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wy_Rw2RLXmMm",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "b321496c-e453-4c1d-9510-4d0e4d797126"
      },
      "source": [
        "import os\n",
        "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
        "import tensorflow as tf\n",
        "print(f'Tensorflow version {tf.__version__}')\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras.layers import Dense, Conv1D, GlobalAveragePooling1D, Embedding\n",
        "import tensorflow_datasets as tfds\n",
        "from tensorflow.keras.models import Model\n",
        "\n",
        "(train_data, test_data), info = tfds.load('imdb_reviews/subwords8k',\n",
        "                                          split=[tfds.Split.TRAIN, tfds.Split.TEST],\n",
        "                                          as_supervised=True, with_info=True)\n",
        "#print(train_data)\n",
        "padded_shapes = ([9000], ())#None.  # this line throws tracing error as the shape of text is varying for each step in an epoch.\n",
        "# as the data size is varying, tf.function will start retracing it\n",
        "# For the demonstration, I used 9000 as max length, but please change it accordingly \n",
        "\n",
        "train_dataset = train_data.shuffle(25000).\\\n",
        "    padded_batch(padded_shapes=padded_shapes, batch_size=16)\n",
        "test_dataset = test_data.shuffle(25000).\\\n",
        "    padded_batch(padded_shapes=padded_shapes, batch_size=16)\n",
        "\n",
        "n_words = info.features['text'].encoder.vocab_size\n",
        "\n",
        "\n",
        "class ConvModel(Model):\n",
        "    def __init__(self):\n",
        "        super(ConvModel, self).__init__()\n",
        "        self.embe = Embedding(n_words, output_dim=16)\n",
        "        self.conv = Conv1D(32, kernel_size=6, activation='elu')\n",
        "        self.glob = GlobalAveragePooling1D()\n",
        "        self.dens = Dense(2)\n",
        "\n",
        "    def call(self, x, training=None, mask=None):\n",
        "        x = self.embe(x)\n",
        "        x = self.conv(x)\n",
        "        x = self.glob(x)\n",
        "        x = self.dens(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "conv = ConvModel()\n",
        "\n",
        "conv(next(iter(train_dataset))[0])\n",
        "\n",
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "\n",
        "train_loss = tf.keras.metrics.Mean()\n",
        "test_loss = tf.keras.metrics.Mean()\n",
        "\n",
        "train_acc = tf.keras.metrics.CategoricalAccuracy()\n",
        "test_acc = tf.keras.metrics.CategoricalAccuracy()\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def train_step(inputs, labels):\n",
        "    with tf.GradientTape() as tape:\n",
        "        logits = conv(inputs, training=True)\n",
        "        loss = loss_object(labels, logits)\n",
        "        train_loss(loss)\n",
        "        train_acc(logits, labels)\n",
        "\n",
        "    gradients = tape.gradient(loss, conv.trainable_variables)\n",
        "    optimizer.apply_gradients(zip(gradients, conv.trainable_variables))\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def test_step(inputs, labels):\n",
        "    logits = conv(inputs, training=False)\n",
        "    loss = loss_object(labels, logits)\n",
        "    test_loss(loss)\n",
        "    test_acc(logits, labels)\n",
        "\n",
        "\n",
        "def learn():\n",
        "    train_loss.reset_states()\n",
        "    test_loss.reset_states()\n",
        "    train_acc.reset_states()\n",
        "    test_acc.reset_states()\n",
        "\n",
        "    for text, target in train_dataset:\n",
        "      #print(text)\n",
        "      #print(target)\n",
        "      train_step(inputs=text, labels=target)\n",
        "\n",
        "    for text, target in test_dataset:\n",
        "        test_step(inputs=text, labels=target)\n",
        "\n",
        "\n",
        "def main(epochs=2):\n",
        "    for epoch in tf.range(1, epochs + 1):\n",
        "        learn()\n",
        "        template = 'TRAIN LOSS {:>5.3f} TRAIN ACC {:.2f} TEST LOSS {:>5.3f} TEST ACC {:.2f}'\n",
        "\n",
        "        print(template.format(\n",
        "            train_loss.result(),\n",
        "            train_acc.result(),\n",
        "            test_loss.result(),\n",
        "            test_acc.result()\n",
        "        ))\n",
        "\n",
        "if __name__ == '__main__':\n",
        "  main(epochs=1)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Tensorflow version 2.2.0-dev20200506\n",
            "TRAIN LOSS 0.679 TRAIN ACC 0.37 TEST LOSS 0.575 TEST ACC 0.41\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bSQDZW2Zpe-P",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}