path: root/tensorflow/contrib/eager/python/examples/workshop/3_inspecting.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Debugging \"graph-first\" models with eager execution",
      "version": "0.3.2",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "[View in Colaboratory](https://colab.research.google.com/gist/alextp/9568ab40f6ed6f9a3ba4736f6aef6127/debugging-graph-first-models-with-eager-execution.ipynb)"
      ]
    },
    {
      "metadata": {
        "id": "mm-t0GuIu1Dt",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "This colab uses eager execution and the Python debugger to modify the execution of a translation model. This combination lets you quickly explore counterfactuals when researching and designing modifications to a model.\n",
        "\n",
        "The model, Transformer from [Tensor2Tensor](https://github.com/tensorflow/tensor2tensor), was originally written with graph building in mind. Executing it eagerly can still be helpful!"
      ]
    },
    {
      "metadata": {
        "id": "gxb1DvIDg4sv",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#@title License (double click to show)\n",
        "# Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\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."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Gx3HA9N1ui64",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 37
        },
        "outputId": "f6986f34-f3e1-44e1-c902-2eb33081acad"
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import pdb\n",
        "tfe = tf.contrib.eager\n",
        "\n",
        "tf.enable_eager_execution()"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3LkOm2ct-Lmc",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 37
        },
        "outputId": "2edc74d9-6bc0-4e78-ab4e-83bf96099ef4"
      },
      "cell_type": "code",
      "source": [
        "!pip install -q -U tensor2tensor\n",
        "from tensor2tensor.models import transformer"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1Z3oMsqV0zB6",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        },
        "outputId": "0a8186ee-c688-457f-c9f6-9a6c1477a93b"
      },
      "cell_type": "code",
      "source": [
        "#@title Create a tensor2tensor translation model, fetch a checkpoint (double click to show)\n",
        "from tensor2tensor import problems\n",
        "from tensor2tensor.utils import trainer_lib\n",
        "from tensor2tensor.utils import registry\n",
        "\n",
        "import numpy as np\n",
        "import os\n",
        "\n",
        "# Setup some directories\n",
        "data_dir = os.path.expanduser(\"~/t2t/data\")\n",
        "tmp_dir = os.path.expanduser(\"~/t2t/tmp\")\n",
        "train_dir = os.path.expanduser(\"~/t2t/train\")\n",
        "checkpoint_dir = os.path.expanduser(\"~/t2t/checkpoints\")\n",
        "tf.gfile.MakeDirs(data_dir)\n",
        "tf.gfile.MakeDirs(tmp_dir)\n",
        "tf.gfile.MakeDirs(train_dir)\n",
        "tf.gfile.MakeDirs(checkpoint_dir)\n",
        "gs_data_dir = \"gs://tensor2tensor-data\"\n",
        "gs_ckpt_dir = \"gs://tensor2tensor-checkpoints/\"\n",
        "\n",
        "# Fetch the problem\n",
        "ende_problem = problems.problem(\"translate_ende_wmt32k\")\n",
        "\n",
        "# Copy the vocab file locally so we can encode inputs and decode model outputs\n",
        "# All vocabs are stored on GCS\n",
        "vocab_name = \"vocab.ende.32768\"\n",
        "vocab_file = os.path.join(gs_data_dir, vocab_name)\n",
        "!gsutil cp {vocab_file} {data_dir}\n",
        "\n",
        "# Get the encoders from the problem\n",
        "encoders = ende_problem.feature_encoders(data_dir)\n",
        "\n",
        "# Setup helper functions for encoding and decoding\n",
        "def encode(input_str, output_str=None):\n",
        "  \"\"\"Input str to features dict, ready for inference\"\"\"\n",
        "  inputs = encoders[\"inputs\"].encode(input_str) + [1]  # add EOS id\n",
        "  batch_inputs = tf.reshape(inputs, [1, -1, 1])  # Make it 3D.\n",
        "  return {\"inputs\": batch_inputs}\n",
        "\n",
        "def decode(integers):\n",
        "  \"\"\"List of ints to str\"\"\"\n",
        "  integers = list(np.squeeze(integers))\n",
        "  if 1 in integers:\n",
        "    integers = integers[:integers.index(1)]\n",
        "  return encoders[\"inputs\"].decode(np.squeeze(integers))\n",
        "\n",
        "# Copy the pretrained checkpoint locally\n",
        "ckpt_name = \"transformer_ende_test\"\n",
        "gs_ckpt = os.path.join(gs_ckpt_dir, ckpt_name)\n",
        "!gsutil -q cp -R {gs_ckpt} {checkpoint_dir}\n",
        "checkpoint_path = tf.train.latest_checkpoint(\n",
        "    os.path.join(checkpoint_dir, ckpt_name))\n",
        "\n",
        "# Create hparams and the model\n",
        "model_name = \"transformer\"\n",
        "hparams_set = \"transformer_base\"\n",
        "\n",
        "hparams = trainer_lib.create_hparams(hparams_set, data_dir=data_dir, problem_name=\"translate_ende_wmt32k\")\n",
        "\n",
        "# NOTE: Only create the model once when restoring from a checkpoint; it's a\n",
        "# Layer and so subsequent instantiations will have different variable scopes\n",
        "# that will not match the checkpoint.\n",
        "translate_model = registry.model(model_name)(hparams, tf.estimator.ModeKeys.EVAL)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying gs://tensor2tensor-data/vocab.ende.32768...\n",
            "/ [1 files][316.4 KiB/316.4 KiB]                                                \n",
            "Operation completed over 1 objects/316.4 KiB.                                    \n",
            "INFO:tensorflow:Setting T2TModel mode to 'eval'\n",
            "INFO:tensorflow:Setting hparams.layer_prepostprocess_dropout to 0.0\n",
            "INFO:tensorflow:Setting hparams.symbol_dropout to 0.0\n",
            "INFO:tensorflow:Setting hparams.attention_dropout to 0.0\n",
            "INFO:tensorflow:Setting hparams.dropout to 0.0\n",
            "INFO:tensorflow:Setting hparams.relu_dropout to 0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "4IblPXLGjuCl",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We've created a Transformer model and fetched an existing training checkpoint. It hasn't created variables yet, and we want to load them from the checkpoint before they're used (restore-on-create) so the first run of the model outputs the correct value. The `tfe.restore_variables_on_create` API looks up variables by name on creation and restores their values."
      ]
    },
    {
      "metadata": {
        "id": "o3MWxcAqJoqG",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "fbc1b1bf-ffbe-4621-b3cb-5eb855fec3a8"
      },
      "cell_type": "code",
      "source": [
        "with tfe.restore_variables_on_create(checkpoint_path):\n",
        "  model_output = translate_model.infer(encode(\"Eager execution\"))\n",
        "print(decode(model_output[\"outputs\"]))"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Greedy Decoding\n",
            "Hinrichtung\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "xk5HV9Hhu9zO",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Using global variable names can get somewhat fragile, so for new code we recommend the object-based `tf.keras.Model.save_weights` or `tf.train.Checkpoint`. However, these require some small code changes to work with existing graph building code.\n",
        "\n",
        "The Transformer model translates \"Eager execution\" in English to \"Hinrichtung\" in German, which refers to capital punishment rather than getting things done. Transformer first encodes the English, then decodes to German. We'll add a debugging hook at the start of the decode phase (once the encodings have been finalized) and see if we can correct the translation."
      ]
    },
    {
      "metadata": {
        "id": "GUGwbYvXZ9-7",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "previous_fast_decode = transformer.fast_decode\n",
        "def debug_fn(*args, **kwargs):\n",
        "  pdb.set_trace()\n",
        "  return previous_fast_decode(*args, **kwargs)  # \"step\" in pdb to step in\n",
        "transformer.fast_decode = debug_fn  # Add our debugging hook to Transformer"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "f61HlvECxJn0",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now that we've \"monkey patched\" the model, we'll drop into a debugger just before decoding starts. In most cases it'd be simpler to add the `pdb.set_trace()` call to the code directly, but in this case we're working with prepackaged library code.\n",
        "\n",
        "First, let's find an encoding which represents the correct sense of \"execution\". Then we'll patch part of that encoding into the encoding of \"Eager execution\" to fix the translation. Feel free to poke around with the debugger (e.g. print a Tensor's value), but your main task is to save the encodings by assigning them to an attribute of the function:\n",
        "\n",
        "```\n",
        "(running the next cell drops you into a pdb shell)\n",
        "step\n",
        "fast_decode.previous_encoding = encoder_output\n",
        "continue\n",
        "\n",
        "```\n",
        "\n",
        "You can type `next` (or `n`) a few times before `continue` to watch the decoding ops run."
      ]
    },
    {
      "metadata": {
        "id": "dX4CPOGSpZrb",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 179
        },
        "outputId": "6de38c31-836f-40ef-b701-e42908172619"
      },
      "cell_type": "code",
      "source": [
        "model_output = translate_model.infer(encode(\"Immediate running\"))\n",
        "print(decode(model_output[\"outputs\"]))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "> <ipython-input-6-ee9b4225ba2a>(4)debug_fn()\n",
            "-> return previous_fast_decode(*args, **kwargs)  # \"step\" in pdb to step in\n",
            "(Pdb) step\n",
            "--Call--\n",
            "> /usr/local/lib/python2.7/dist-packages/tensor2tensor/models/transformer.py(427)fast_decode()\n",
            "-> def fast_decode(encoder_output,\n",
            "(Pdb) fast_decode.previous_encoding = encoder_output\n",
            "(Pdb) continue\n",
            "Sofortige Durchführung\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "-ZEZciV4FpLo",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now we have an encoding saved which gets the correct sense for \"execution\"."
      ]
    },
    {
      "metadata": {
        "id": "QeC_oDVqHD_v",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 179
        },
        "outputId": "253c9af1-003e-46bd-8bf5-db968cf6a8cf"
      },
      "cell_type": "code",
      "source": [
        "# Assumes you followed the pdb instructions above!\n",
        "transformer.fast_decode.previous_encoding"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=9528, shape=(1, 4, 512), dtype=float32, numpy=\n",
              "array([[[-0.15239455,  0.12273102, -0.11209048, ..., -0.12478986,\n",
              "          0.37216735, -0.40987235],\n",
              "        [-0.2686283 ,  0.51448774,  0.03650613, ...,  0.08731575,\n",
              "          0.51110077, -0.6646815 ],\n",
              "        [-0.24441548,  0.36622533,  0.11685672, ...,  0.21941349,\n",
              "         -0.03304008, -0.579611  ],\n",
              "        [-0.03339856, -0.01185844,  0.00579634, ...,  0.00294734,\n",
              "          0.00136655, -0.01362935]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "metadata": {
        "id": "bC9JjeDcHEav",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Let's replace part of the encoding for \"Eager execution\" with the encoding of \"Immediate running\".\n",
        "\n",
        "Again we'll drop into a pdb shell. This time we'll run some TensorFlow operations to patch the encodings while the model is running.\n",
        "\n",
        "```\n",
        "(running the next cell again drops you into a pdb shell)\n",
        "step\n",
        "encoder_output = tf.concat([fast_decode.previous_encoding[:, :3], encoder_output[:, 3:]], axis=1)\n",
        "continue\n",
        "```"
      ]
    },
    {
      "metadata": {
        "id": "t2as_Kn1h65G",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 179
        },
        "outputId": "5b4e546e-3bb4-4761-c545-467b631e3ffe"
      },
      "cell_type": "code",
      "source": [
        "model_output = translate_model.infer(encode(\"Eager execution\"))\n",
        "print(decode(model_output[\"outputs\"]))"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "> <ipython-input-6-ee9b4225ba2a>(4)debug_fn()\n",
            "-> return previous_fast_decode(*args, **kwargs)  # \"step\" in pdb to step in\n",
            "(Pdb) step\n",
            "--Call--\n",
            "> /usr/local/lib/python2.7/dist-packages/tensor2tensor/models/transformer.py(427)fast_decode()\n",
            "-> def fast_decode(encoder_output,\n",
            "(Pdb) encoder_output = tf.concat([fast_decode.previous_encoding[:, :3], encoder_output[:, 3:]], axis=1)\n",
            "(Pdb) continue\n",
            "sofortige Ausführung\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "rK6tYZ23I2cm",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We get a different decoding, with the correct sense of \"execution\". Likely we're keeping just the encoding of \"tion\" from \"Eager execution\", so no great breakthrough in translation modeling.\n",
        "\n",
        "Similarly it's possible to modify attention vectors, or change words during decoding to help debug a beam search."
      ]
    },
    {
      "metadata": {
        "id": "Nb-4ipYNRWxA",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "This colab was adapted from the [Tensor2Tensor colab](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb). Credit to Ankur Taly for its concept."
      ]
    }
  ]
}