{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "FYLyuStTYesc"
      },
      "outputs": [],
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "form",
        "id": "PVm-iEoxYesf"
      },
      "outputs": [],
      "source": [
        "#@title 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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3MPf91rVYesq"
      },
      "source": [
        "# 使用 Grappler 优化 TensorFlow 计算图"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zmNCsZlgYesr"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://tensorflow.google.cn/guide/graph_optimization\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\">在 TensorFlow.org 上查看</a>\n",
        "</td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/graph_optimization.ipynb\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\">在 Google Colab 中运行</a></td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/graph_optimization.ipynb\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\">在 Github 上查看源代码</a>\n",
        "</td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/guide/graph_optimization.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\">下载笔记本</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/media/pc/data/lxw/ai/d2py/doc/libs/tf-chaos/guide\n"
          ]
        }
      ],
      "source": [
        "%cd ..\n",
        "from set_env import temp_dir"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l0qacLgyYess"
      },
      "source": [
        "## 概述\n",
        "\n",
        "TensorFlow 同时使用计算图和 Eager Execution 来执行计算。一个 `tf.Graph` 包含一组代表计算单元的 `tf.Operation` 对象（运算）和一组代表在运算之间流动的数据单元的 `tf.Tensor` 对象。\n",
        "\n",
        "Grappler 是 TensorFlow 运行时中的默认计算图优化系统。Grappler 通过计算图简化和其他高级优化（例如利用内嵌函数体实现程序间优化），在计算图模式（在 `tf.function` 内）下应用优化以提高 TensorFlow 计算的性能。优化 `tf.Graph` 还可以通过优化计算图节点到计算资源的映射来减少设备峰值内存使用量并提高硬件利用率。\n",
        "\n",
        "使用 `tf.config.optimizer.set_experimental_options()` 可以更好地控制 `tf.Graph` 优化。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-zkJgR5Yesw"
      },
      "source": [
        "## 可用的计算图优化器\n",
        "\n",
        "Grappler 通过称为 `MetaOptimizer` 的顶级驱动程序执行计算图优化。TensorFlow 提供以下计算图优化器：\n",
        "\n",
        "- *常量折叠优化器* - 通过折叠计算图中的常量节点来静态推断张量的值（如可能），并使用常量使结果具体化。\n",
        "- *算术优化器* - 通过消除常见的子表达式并简化算术语句来简化算术运算。\n",
        "- *布局优化器 -* 优化张量布局以更高效地执行依赖于数据格式的运算，例如卷积。\n",
        "- *重新映射优化器* - 通过将常见的子计算图替换为经过优化的融合一体化内核，将子计算图重新映射到更高效的实现上。\n",
        "- *内存优化器 -* 分析计算图以检查每个运算的峰值内存使用量，并插入 CPU-GPU 内存复制操作以将 GPU 内存交换到 CPU，从而减少峰值内存使用量。\n",
        "- *依赖项优化器 -* 移除或重新排列控制依赖项，以缩短模型步骤的关键路径或实现其他优化。另外，还移除了实际上是无运算的节点，例如 Identity。\n",
        "- *剪枝优化器 -* 修剪对计算图的输出没有影响的节点。通常会首先运行剪枝来减小计算图的大小并加快其他 Grappler 传递中的处理速度。\n",
        "- *函数优化器 -* 优化 TensorFlow 程序的函数库，并内嵌函数体以实现其他程序间优化。\n",
        "- *形状优化器 -* 优化对形状和形状相关信息进行运算的子计算图。\n",
        "- *自动并行优化器 -* 通过沿批次维度拆分来自动并行化计算图。默认情况下，此优化器处于关闭状态。\n",
        "- *循环优化器 -* 通过将循环不变式子计算图提升到循环外并通过移除循环中的冗余堆栈运算来优化计算图控制流。另外，还优化具有静态已知行程计数的循环，并移除条件语句中静态已知的无效分支。\n",
        "- *范围分配器优化器 -* 引入范围分配器以减少数据移动并合并某些运算。\n",
        "- *固定到主机优化器 -* 将小型运算交换到 CPU 上。默认情况下，此优化器处于关闭状态。\n",
        "- *自动混合精度优化器 -* 在适用的情况下将数据类型转换为 float16 以提高性能。目前仅适用于 GPU。\n",
        "- *调试剥离器 -* 从计算图中剥离与调试运算相关的节点，例如 `tf.debugging.Assert`、`tf.debugging.check_numerics` 和 `tf.print`。默认情况下，此优化器处于关闭状态。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZAUsxyWYess"
      },
      "source": [
        "## 设置\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "6BRIDzO6ypoY"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import timeit\n",
        "import traceback\n",
        "import contextlib\n",
        "\n",
        "\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1O-XL1nxJX0X"
      },
      "source": [
        "创建上下文管理器以轻松切换优化器状态。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "uRuhVoAlYesz"
      },
      "outputs": [],
      "source": [
        "@contextlib.contextmanager\n",
        "def options(options):\n",
        "  old_opts = tf.config.optimizer.get_experimental_options()\n",
        "  tf.config.optimizer.set_experimental_options(options)\n",
        "  try:\n",
        "    yield\n",
        "  finally:\n",
        "    tf.config.optimizer.set_experimental_options(old_opts)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E2o4kZtK0DoA"
      },
      "source": [
        "## 比较使用和不使用 Grappler 时的执行性能\n",
        "\n",
        "TensorFlow 2 及更高版本默认情况下会以 [Eager](../eager.md) 模式执行。使用 `tf.function` 可将默认执行切换为“计算图”模式。Grappler 在后台自动运行，以应用上述计算图优化并提高执行性能。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3sh8RoLJ96IT"
      },
      "source": [
        "### 常量折叠优化器\n",
        "\n",
        "作为一个初步的示例，考虑一个对常量执行运算并返回输出的函数。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "jOW_OSzMJEvN"
      },
      "outputs": [],
      "source": [
        "def test_function_1():\n",
        "  @tf.function\n",
        "  def simple_function(input_arg):\n",
        "    print('Tracing!')\n",
        "    a = tf.constant(np.random.randn(2000,2000), dtype = tf.float32)\n",
        "    c = a\n",
        "    for n in range(50):\n",
        "      c = c@a\n",
        "    return tf.reduce_mean(c+input_arg)\n",
        "\n",
        "  return simple_function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFVgUhhzLKIo"
      },
      "source": [
        "关闭常量折叠优化器并执行以下函数："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "KDMGsOHrJqKD"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'constant_folding': False, 'disable_model_pruning': False, 'disable_meta_optimizer': False}\n",
            "Tracing!\n",
            "Vanilla execution: 0.0024620250333100557 s\n"
          ]
        }
      ],
      "source": [
        "with options({'constant_folding': False}):\n",
        "  print(tf.config.optimizer.get_experimental_options())\n",
        "  simple_function = test_function_1()\n",
        "  # Trace once\n",
        "  x = tf.constant(2.2)\n",
        "  simple_function(x)\n",
        "  print(\"Vanilla execution:\", timeit.timeit(lambda: simple_function(x), number = 1), \"s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykMXfo8qO41z"
      },
      "source": [
        "启用常量折叠优化器，然后再次执行函数以观察函数执行的加速情况。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "73pc0gfFKY8l"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'constant_folding': True, 'disable_model_pruning': False, 'disable_meta_optimizer': False}\n",
            "Tracing!\n",
            "Constant folded execution: 0.0006806210149079561 s\n"
          ]
        }
      ],
      "source": [
        "with options({'constant_folding': True}):\n",
        "  print(tf.config.optimizer.get_experimental_options())\n",
        "  simple_function = test_function_1()\n",
        "  # Trace once\n",
        "  x = tf.constant(2.2)\n",
        "  simple_function(x)\n",
        "  print(\"Constant folded execution:\", timeit.timeit(lambda: simple_function(x), number = 1), \"s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83w8rfcRVhWb"
      },
      "source": [
        "### 调试剥离器优化器\n",
        "\n",
        "考虑一个检查其输入参数的数值并返回自身的简单函数。 "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "j2DvqEr8haut"
      },
      "outputs": [],
      "source": [
        "def test_function_2():\n",
        "  @tf.function\n",
        "  def simple_func(input_arg):\n",
        "    output = input_arg\n",
        "    tf.debugging.check_numerics(output, \"Bad!\")\n",
        "    return output\n",
        "  return simple_func"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ywKG3WRbpYB8"
      },
      "source": [
        "首先，在调试剥离器优化器关闭的情况下执行该函数。 "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "LsE-y6iQWSwH"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Traceback (most recent call last):\n",
            "  File \"/tmp/ipykernel_4128893/3616845043.py\", line 4, in <module>\n",
            "    test_func(p1)\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/tensorflow/python/util/traceback_utils.py\", line 153, in error_handler\n",
            "    raise e.with_traceback(filtered_tb) from None\n",
            "tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:\n",
            "\n",
            "Detected at node CheckNumerics defined at (most recent call last):\n",
            "  File \"<frozen runpy>\", line 198, in _run_module_as_main\n",
            "\n",
            "  File \"<frozen runpy>\", line 88, in _run_code\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel_launcher.py\", line 17, in <module>\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/traitlets/config/application.py\", line 1075, in launch_instance\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelapp.py\", line 701, in start\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/tornado/platform/asyncio.py\", line 205, in start\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/asyncio/base_events.py\", line 639, in run_forever\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/asyncio/base_events.py\", line 1985, in _run_once\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/asyncio/events.py\", line 88, in _run\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py\", line 534, in dispatch_queue\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py\", line 523, in process_one\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py\", line 429, in dispatch_shell\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/kernelbase.py\", line 767, in execute_request\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/ipkernel.py\", line 429, in do_execute\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/ipykernel/zmqshell.py\", line 549, in run_cell\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py\", line 3075, in run_cell\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py\", line 3130, in _run_cell\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/async_helpers.py\", line 128, in _pseudo_sync_runner\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py\", line 3334, in run_cell_async\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py\", line 3517, in run_ast_nodes\n",
            "\n",
            "  File \"/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/IPython/core/interactiveshell.py\", line 3577, in run_code\n",
            "\n",
            "  File \"/tmp/ipykernel_4128893/3616845043.py\", line 4, in <module>\n",
            "\n",
            "  File \"/tmp/ipykernel_4128893/2241890286.py\", line 5, in simple_func\n",
            "\n",
            "Bad! : Tensor had Inf values\n",
            "\t [[{{node CheckNumerics}}]] [Op:__inference_simple_func_128]\n"
          ]
        }
      ],
      "source": [
        "test_func = test_function_2()\n",
        "p1 = tf.constant(float('inf'))\n",
        "try:\n",
        "  test_func(p1)\n",
        "except tf.errors.InvalidArgumentError as e:\n",
        "  traceback.print_exc(limit=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "URHpboM8xLN6"
      },
      "source": [
        "由于 `test_func` 的 `Inf` 参数，`tf.debugging.check_numerics` 引发了参数无效错误。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuPSha9YmJRo"
      },
      "source": [
        "启用调试剥离器优化器，然后再次执行该函数。 "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "UPJ7ygHnWP6B"
      },
      "outputs": [],
      "source": [
        "with options({'debug_stripper': True}):\n",
        "  test_func2 = test_function_2()\n",
        "  p1 = tf.constant(float('inf'))\n",
        "  try:\n",
        "    test_func2(p1)\n",
        "  except tf.errors.InvalidArgumentError as e:\n",
        "    traceback.print_exc(limit=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nAsESNUB1QpI"
      },
      "source": [
        "调试剥离器优化器从计算图中剥离 `tf.debug.check_numerics` 节点并执行该函数，而不会引发任何错误。 "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wHC6tR9GvFgW"
      },
      "source": [
        "## 总结\n",
        "\n",
        "TensorFlow 运行时会在执行之前使用 Grappler 自动优化计算图。使用 `tf.config.optimizer.set_experimental_options` 可启用或停用各个计算图优化器。\n",
        "\n",
        "有关 Grappler 的更多信息，请参阅 <a class=\"external\" href=\"http://web.stanford.edu/class/cs245/slides/TFGraphOptimizationsStanford.pdf\">TensorFlow 计算图优化</a>。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "graph_optimization.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "xxx",
      "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.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}