site/zh-cn/guide/migrate/migrating_checkpoints.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6bYaCABobL5q"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "FlUw7tSKbtg4"
      },
      "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": "61dp4Hg5ksTC"
      },
      "source": [
        "# 迁移模型检查点\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://tensorflow.google.cn/guide/migrate/migrating_checkpoints\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\">在 TensorFlow.org上查看</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/migrate/migrating_checkpoints.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/migrate/migrating_checkpoints.ipynb\">     <img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\">     在 GitHub 上查看源代码</a></td>\n",
        "  <td>     <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/guide/migrate/migrating_checkpoints.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\">下载笔记本</a>   </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "avuMwzscPnHh"
      },
      "source": [
        "注：使用 `tf.compat.v1.Saver` 保存的检查点通常称为 *TF1 或基于名称的*检查点。使用 `tf.train.Checkpoint` 保存的检查点称为 *TF2 或基于对象的*检查点。\n",
        "\n",
        "## 概述\n",
        "\n",
        "本指南假定您有一个使用 [`tf.compat.v1.Saver`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/Saver) 保存和加载检查点的模型，并且想要使用 TF2 [`tf.train.Checkpoint`](https://tensorflow.google.cn/api_docs/python/tf/train/Checkpoint) API 迁移代码，或者使用 TF2 模型中既有的检查点。\n",
        "\n",
        "下面是可能会遇到的一些常见情形：\n",
        "\n",
        "**情形 1**\n",
        "\n",
        "以前的训练运行中存在现有的 TF1 检查点，需要加载或转换为 TF2。\n",
        "\n",
        "- 要在 TF2 中加载 TF1 检查点，请参阅[*在 TF2 中加载 TF1 检查点*](#load-tf1-in-tf2)代码段。\n",
        "- 要将检查点转换为 TF2，请参阅[*检查点转换*](#checkpoint-conversion)。\n",
        "\n",
        "**情形 2**\n",
        "\n",
        "您正在以一种存在更改变量名称和路径的风险的方式调整您的模型（例如，从 `get_variable` 增量迁移到显式 `tf.Variable` 创建时），并且希望在此过程中保持现有检查点的保存/加载。\n",
        "\n",
        "请参阅[*如何在模型迁移期间保持检查点兼容性*](#maintain-checkpoint-compat)部分\n",
        "\n",
        "**情形 3**\n",
        "\n",
        "您正在将训练代码和检查点迁移到 TF2，但您的推断流水线目前仍需要 TF1 检查点（为了生产稳定性）。\n",
        "\n",
        "*选项 1*\n",
        "\n",
        "训练时同时保存 TF1 和 TF2 检查点。\n",
        "\n",
        "- 请参阅[*在 TF2 中保存 TF1 检查点*](#save-tf1-in-tf2)\n",
        "\n",
        "*选项 2*\n",
        "\n",
        "将 TF2 检查点转换为 TF1。\n",
        "\n",
        "- 请参阅[*检查点转换*](#checkpoint-conversion)\n",
        "\n",
        "---\n",
        "\n",
        "下面的示例显示了 TF1/TF2 中保存和加载检查点的所有组合，因此可以灵活地确定如何迁移模型。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TaYgaekzOAHf"
      },
      "source": [
        "## 安装"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kcvTd5QhZ78L"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1\n",
        "\n",
        "def print_checkpoint(save_path):\n",
        "  reader = tf.train.load_checkpoint(save_path)\n",
        "  shapes = reader.get_variable_to_shape_map()\n",
        "  dtypes = reader.get_variable_to_dtype_map()\n",
        "  print(f\"Checkpoint at '{save_path}':\")\n",
        "  for key in shapes:\n",
        "    print(f\"  (key='{key}', shape={shapes[key]}, dtype={dtypes[key].name}, \"\n",
        "          f\"value={reader.get_tensor(key)})\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gO8Q6QkulJlj"
      },
      "source": [
        "## 从 TF1 到 TF2 的变化\n",
        "\n",
        "如果您对 TF1 和 TF2 之间发生了哪些变化以及我们所说的“基于名称”(TF1) 与“基于对象”(TF2) 的检查点的含义感到好奇，请阅读此部分。\n",
        "\n",
        "这两种类型的检查点实际上以相同的格式保存，本质上是一个键值表。不同之处在于键的生成方式。\n",
        "\n",
        "基于名称的检查点中的键是**变量的名称**。基于对象的检查点中的键指向**从根对象到变量的路径**（下面的示例将有助于更好地理解这段话的含义）。\n",
        "\n",
        "首先，保存一些检查点：\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8YXzbXvOWvdF"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  c = tf1.get_variable('scoped/c', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  with tf1.Session() as sess:\n",
        "    saver = tf1.train.Saver()\n",
        "    sess.run(a.assign(1))\n",
        "    sess.run(b.assign(2))\n",
        "    sess.run(c.assign(3))\n",
        "    saver.save(sess, 'tf1-ckpt')\n",
        "\n",
        "print_checkpoint('tf1-ckpt')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "raOych1UaJzl"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(5.0, name='a')\n",
        "b = tf.Variable(6.0, name='b')\n",
        "with tf.name_scope('scoped'):\n",
        "  c = tf.Variable(7.0, name='c')\n",
        "\n",
        "ckpt = tf.train.Checkpoint(variables=[a, b, c])\n",
        "save_path_v2 = ckpt.save('tf2-ckpt')\n",
        "print_checkpoint(save_path_v2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYyLhTYszcpl"
      },
      "source": [
        "如果您查看 `tf2-ckpt` 中的键，它们全部指向每个变量的对象路径。例如，变量 `a` 是 `variables` 列表中的第一个元素，因此它的键变为 `variables/0/...`（请尽管忽略 .ATTRIBUTES/VARIABLE_VALUE 常量）。\n",
        "\n",
        "仔细检查下面的 `Checkpoint` 对象："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kLOxvoosg4Al"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(0.)\n",
        "b = tf.Variable(0.)\n",
        "c = tf.Variable(0.)\n",
        "root = ckpt = tf.train.Checkpoint(variables=[a, b, c])\n",
        "print(\"root type =\", type(root).__name__)\n",
        "print(\"root.variables =\", root.variables)\n",
        "print(\"root.variables[0] =\", root.variables[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Qaed1yAm3Ar"
      },
      "source": [
        "尝试使用下面的代码段，看看检查点键如何随对象结构变化："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EdHJXlZOyDnn"
      },
      "outputs": [],
      "source": [
        "module = tf.Module()\n",
        "module.d = tf.Variable(0.)\n",
        "test_ckpt = tf.train.Checkpoint(v={'a': a, 'b': b}, \n",
        "                                c=c,\n",
        "                                module=module)\n",
        "test_ckpt_path = test_ckpt.save('root-tf2-ckpt')\n",
        "print_checkpoint(test_ckpt_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8iWitEsayDWs"
      },
      "source": [
        "*为什么 TF2 使用这种机制？*\n",
        "\n",
        "TF2 中没有更多的全局计算图，因此变量名称是不可靠的，并且程序之间可能存在不一致。TF2 鼓励使用面向对象的建模方法，其中变量归层所有，层归模型所有：\n",
        "\n",
        "```\n",
        "variable = tf.Variable(...)\n",
        "layer.variable_name = variable\n",
        "model.layer_name = layer\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9kv9SmyVjGLA"
      },
      "source": [
        "## 如何在模型迁移期间保持检查点兼容性\n",
        "\n",
        "<a name=\"maintain-checkpoint-compat\"></a>\n",
        "\n",
        "迁移过程中的一个重要步骤是*确保所有变量都被初始化为正确的值*，这反过来又允许您验证运算/函数是否正在执行正确的计算。为此，您必须考虑迁移各个阶段中模型之间的**检查点兼容性**。本质上，本部分回答了一个问题，即*如何在更改模型时继续使用相同的检查点*。\n",
        "\n",
        "为了提高灵活性，下面是维护检查点兼容性的三种方法：\n",
        "\n",
        "1. 模型的**变量名称和之前相同**。\n",
        "2. 模型有不同的变量名称，并维护一个将检查点中的变量名称映射到新名称的**分配映射**。\n",
        "3. 模型有不同的变量名称，并维护了一个存储所有变量的 **TF2 检查点对象**。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L5JhCyPZDx43"
      },
      "source": [
        "### 当变量名称匹配时\n",
        "\n",
        "长标题：如何在变量名称匹配时重用检查点。\n",
        "\n",
        "短答案：可以使用 `tf1.train.Saver` 或 `tf.train.Checkpoint` 直接加载既有的检查点。\n",
        "\n",
        "---\n",
        "\n",
        "如果您使用的是 `tf.compat.v1.keras.utils.track_tf1_style_variables`，那么它将确保模型变量名称与以前相同。您还可以手动确保变量名称匹配。\n",
        "\n",
        "当迁移模型中的变量名称匹配时，您可以直接使用 `tf.train.Checkpoint` 或 `tf.compat.v1.train.Saver` 加载检查点。这两个 API 都与 Eager 和计算图模式兼容，因此您可以在迁移的任何阶段使用它们。\n",
        "\n",
        "注：您可以使用 `tf.train.Checkpoint` 加载 TF1 检查点，但如果没有复杂的名称匹配，则不能使用 `tf.compat.v1.Saver` 加载 TF2 检查点。\n",
        "\n",
        "下面是对不同模型使用相同检查点的示例。首先，使用 `tf1.train.Saver` 保存一个 TF1 检查点："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijlHS96URsfR"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  c = tf1.get_variable('scoped/c', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  with tf1.Session() as sess:\n",
        "    saver = tf1.train.Saver()\n",
        "    sess.run(a.assign(1))\n",
        "    sess.run(b.assign(2))\n",
        "    sess.run(c.assign(3))\n",
        "    save_path = saver.save(sess, 'tf1-ckpt')\n",
        "print_checkpoint(save_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zg7nWZphQD9u"
      },
      "source": [
        "下面的示例使用 `tf.compat.v1.Saver` 在 Eager 模式下加载检查点："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y4K16m0PPncQ"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(0.0, name='a')\n",
        "b = tf.Variable(0.0, name='b')\n",
        "with tf.name_scope('scoped'):\n",
        "  c = tf.Variable(0.0, name='c')\n",
        "\n",
        "# With the removal of collections in TF2, you must pass in the list of variables\n",
        "# to the Saver object:\n",
        "saver = tf1.train.Saver(var_list=[a, b, c])\n",
        "saver.restore(sess=None, save_path=save_path)\n",
        "print(f\"loaded values of [a, b, c]:  [{a.numpy()}, {b.numpy()}, {c.numpy()}]\")\n",
        "\n",
        "# Saving also works in eager (sess must be None).\n",
        "path = saver.save(sess=None, save_path='tf1-ckpt-saved-in-eager')\n",
        "print_checkpoint(path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWnq1f5yAPkq"
      },
      "source": [
        "下一个代码段使用 TF2 API `tf.train.Checkpoint` 加载检查点："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "StyrzwGvW1YZ"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(0.0, name='a')\n",
        "b = tf.Variable(0.0, name='b')\n",
        "with tf.name_scope('scoped'):\n",
        "  c = tf.Variable(0.0, name='c')\n",
        "\n",
        "# Without the name_scope, name=\"scoped/c\" works too:\n",
        "c_2 = tf.Variable(0.0, name='scoped/c')\n",
        "\n",
        "print(\"Variable names: \")\n",
        "print(f\"  a.name = {a.name}\")\n",
        "print(f\"  b.name = {b.name}\")\n",
        "print(f\"  c.name = {c.name}\")\n",
        "print(f\"  c_2.name = {c_2.name}\")\n",
        "\n",
        "# Restore the values with tf.train.Checkpoint\n",
        "ckpt = tf.train.Checkpoint(variables=[a, b, c, c_2])\n",
        "ckpt.restore(save_path)\n",
        "print(f\"loaded values of [a, b, c, c_2]:  [{a.numpy()}, {b.numpy()}, {c.numpy()}, {c_2.numpy()}]\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DYYgbj8F7Yb7"
      },
      "source": [
        "#### TF2 中的变量名称\n",
        "\n",
        "- 变量仍然具有您可以设置的 `name` 参数。\n",
        "- Keras 模型还采用 `name` 参数，并将其设置为变量的前缀。\n",
        "- `v1.name_scope` 函数可用于设置变量名前缀，这与 `tf.variable_scope` 截然不同。它只影响名称，而不跟踪变量和重用。\n",
        "\n",
        "`tf.compat.v1.keras.utils.track_tf1_style_variables` 装饰器是一个填充码，它通过保持 `tf.variable_scope` 和 `tf.compat.v1.get_variable` 的命名和重用语义不变，帮助您维护变量名称和 TF1 检查点兼容性。如需了解详情，请参阅[模型映射指南](./model_mapping.ipynb)。\n",
        "\n",
        "**注 1：如果您使用填充码，请使用 TF2 API 加载您的检查点（即使使用预训练的 TF1 检查点）。**\n",
        "\n",
        "请参阅*检查点 Keras* 部分。\n",
        "\n",
        "**注 2：从 `get_variable` 迁移到 `tf.Variable` 时：**\n",
        "\n",
        "如果您的填充码装饰层或模块包含一些使用 `tf.Variable` 而不是 `tf.compat.v1.get_variable` 的变量（或 Keras 层/模型），并作为属性附加/以面向对象的方式进行跟踪，则与 Eager Execution 期间相比，它们在 TF1.x 计算图/会话中可能具有不同的变量命名语义。\n",
        "\n",
        "简而言之，在 TF2 中运行时，*这些名称可能不是预期的名称*。\n",
        "\n",
        "警告：变量在 Eager Execution 中可能有重复的名称，如果基于名称的检查点中的多个变量需要映射到相同的名称，这可能会导致问题。可以使用 `tf.name_scope` 和层构造函数或 `tf.Variable` `name` 参数显式调整层和变量名称，从而调整变量名称并确保没有重复。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NkUQJUUyjOJz"
      },
      "source": [
        "### 维护分配映射\n",
        "\n",
        "分配映射通常用于在 TF1 模型之间传递权重，如果变量名称发生变化，也可以在模型迁移期间使用。\n",
        "\n",
        "您可以将这些映射与 [`tf.compat.v1.train.init_from_checkpoint`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/init_from_checkpoint)、[`tf.compat.v1.train.Saver`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/Saver) 和 [`tf.train.load_checkpoint`](https://tensorflow.google.cn/api_docs/python/tf/train/load_checkpoint) 一起使用，以将权重加载到变量或范围名称可能已更改的模型中。\n",
        "\n",
        "本部分中的示例将使用之前保存的检查点："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PItyo7DdJ6Ek"
      },
      "outputs": [],
      "source": [
        "print_checkpoint('tf1-ckpt')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rPryV_WBJrI3"
      },
      "source": [
        "#### 使用 `init_from_checkpoint` 加载\n",
        "\n",
        "[`tf1.train.init_from_checkpoint`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/init_from_checkpoint) 必须在计算图/会话中调用，因为它将值置于变量初始值设定项中，而不是创建分配运算。\n",
        "\n",
        "您可以使用 `assignment_map` 参数来配置变量的加载方式。从文档中：\n",
        "\n",
        "> 分配映射支持以下语法：\n",
        "\n",
        "- `'checkpoint_scope_name/': 'scope_name/'` - 将从 `checkpoint_scope_name` 加载当前 `scope_name` 中的所有变量，并具有匹配的张量名称。\n",
        "- `'checkpoint_scope_name/some_other_variable': 'scope_name/variable_name'` - 将从 `checkpoint_scope_name/some_other_variable` 初始化 `scope_name/variable_name` 变量。\n",
        "- `'scope_variable_name': variable` - 将使用来自检查点的张量 'scope_variable_name' 初始化给定的 `tf.Variable` 对象。\n",
        "- `'scope_variable_name': list(variable)` - 将使用来自检查点的张量 'scope_variable_name' 初始化分区变量列表。\n",
        "- `'/': 'scope_name/'` - 将从检查点的根目录（例如，无范围）加载当前 `scope_name` 中的所有变量。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZM_7OzRpdH0A"
      },
      "outputs": [],
      "source": [
        "# Restoring with tf1.train.init_from_checkpoint:\n",
        "\n",
        "# A new model with a different scope for the variables.\n",
        "with tf.Graph().as_default() as g:\n",
        "  with tf1.variable_scope('new_scope'):\n",
        "    a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "    b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "    c = tf1.get_variable('scoped/c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "  with tf1.Session() as sess:\n",
        "    # The assignment map will remap all variables in the checkpoint to the\n",
        "    # new scope:\n",
        "    tf1.train.init_from_checkpoint(\n",
        "        'tf1-ckpt',\n",
        "        assignment_map={'/': 'new_scope/'})\n",
        "    # `init_from_checkpoint` adds the initializers to these variables.\n",
        "    # Use `sess.run` to run these initializers.\n",
        "    sess.run(tf1.global_variables_initializer())\n",
        "\n",
        "    print(\"Restored [a, b, c]: \", sess.run([a, b, c]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Za_8xhFWKVlH"
      },
      "source": [
        "#### 使用 `tf1.train.Saver` 加载\n",
        "\n",
        "与  `init_from_checkpoint` 不同，[`tf.compat.v1.train.Saver`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/Saver) 同时支持在计算图模式和 Eager 模式下运行。`var_list` 参数可以接受字典，但它必须将变量名称映射到 `tf.Variable` 对象。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IiKNmdGJgoX9"
      },
      "outputs": [],
      "source": [
        "# Restoring with tf1.train.Saver (works in both graph and eager):\n",
        "\n",
        "# A new model with a different scope for the variables.\n",
        "with tf1.variable_scope('new_scope'):\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                      initializer=tf1.zeros_initializer())\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                      initializer=tf1.zeros_initializer())\n",
        "  c = tf1.get_variable('scoped/c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "# Initialize the saver with a dictionary with the original variable names:\n",
        "saver = tf1.train.Saver({'a': a, 'b': b, 'scoped/c': c})\n",
        "saver.restore(sess=None, save_path='tf1-ckpt')\n",
        "print(\"Restored [a, b, c]: \", [a.numpy(), b.numpy(), c.numpy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JsgCXt3Ly-h"
      },
      "source": [
        "#### 使用 `tf.train.load_checkpoint` 加载\n",
        "\n",
        "如果您需要精确控制变量值，则此选项适合您。同样，这适用于计算图和 Eager 模式。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pc39Bh6JMso6"
      },
      "outputs": [],
      "source": [
        "# Restoring with tf.train.load_checkpoint (works in both graph and eager):\n",
        "\n",
        "# A new model with a different scope for the variables.\n",
        "with tf.Graph().as_default() as g:\n",
        "  with tf1.variable_scope('new_scope'):\n",
        "    a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "    b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "    c = tf1.get_variable('scoped/c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "  with tf1.Session() as sess:\n",
        "    # It may be easier writing a loop if your model has a lot of variables.\n",
        "    reader = tf.train.load_checkpoint('tf1-ckpt')\n",
        "    sess.run(a.assign(reader.get_tensor('a')))\n",
        "    sess.run(b.assign(reader.get_tensor('b')))\n",
        "    sess.run(c.assign(reader.get_tensor('scoped/c')))\n",
        "    print(\"Restored [a, b, c]: \", sess.run([a, b, c]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nBSTJVCNDKed"
      },
      "source": [
        "### 维护 TF2 检查点对象\n",
        "\n",
        "如果在迁移过程中变量和范围名称可能会发生很大变化，请使用 `tf.train.Checkpoint` 和 TF2 检查点。TF2 使用**对象结构**而不是变量名（有关详情，请参阅*从 TF1 到 TF2 的变化*）。\n",
        "\n",
        "简而言之，在创建 `tf.train.Checkpoint` 来保存或恢复检查点时，请确保它使用相同的**顺序**（对于列表）和**键**（对于 `Checkpoint` 初始值设定项的字典和关键字参数）。检查点兼容性的一些示例：\n",
        "\n",
        "```\n",
        "ckpt = tf.train.Checkpoint(foo=[var_a, var_b])\n",
        "\n",
        "# compatible with ckpt\n",
        "tf.train.Checkpoint(foo=[var_a, var_b])\n",
        "\n",
        "# not compatible with ckpt\n",
        "tf.train.Checkpoint(foo=[var_b, var_a])\n",
        "tf.train.Checkpoint(bar=[var_a, var_b])\n",
        "```\n",
        "\n",
        "下面的代码示例显示了如何使用“相同”的 `tf.train.Checkpoint` 来加载具有不同名称的变量。首先，保存一个 TF2 检查点：\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tCSkz_-Tbct6"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(1))\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(2))\n",
        "  with tf1.variable_scope('scoped'):\n",
        "    c = tf1.get_variable('c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.constant_initializer(3))\n",
        "  with tf1.Session() as sess:\n",
        "    sess.run(tf1.global_variables_initializer())\n",
        "    print(\"[a, b, c]: \", sess.run([a, b, c]))\n",
        "\n",
        "    # Save a TF2 checkpoint\n",
        "    ckpt = tf.train.Checkpoint(unscoped=[a, b], scoped=[c])\n",
        "    tf2_ckpt_path = ckpt.save('tf2-ckpt')\n",
        "    print_checkpoint(tf2_ckpt_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "62MWdZMxezeP"
      },
      "source": [
        "即使变量/范围名称发生变化，也可以继续使用 `tf.train.Checkpoint`："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vh61SGeqb27b"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a_different_name', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  b = tf1.get_variable('b_different_name', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.zeros_initializer())\n",
        "  with tf1.variable_scope('different_scope'):\n",
        "    c = tf1.get_variable('c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.zeros_initializer())\n",
        "  with tf1.Session() as sess:\n",
        "    sess.run(tf1.global_variables_initializer())\n",
        "    print(\"Initialized [a, b, c]: \", sess.run([a, b, c]))\n",
        "\n",
        "    ckpt = tf.train.Checkpoint(unscoped=[a, b], scoped=[c])\n",
        "    # `assert_consumed` validates that all checkpoint objects are restored from\n",
        "    # the checkpoint. `run_restore_ops` is required when running in a TF1\n",
        "    # session.\n",
        "    ckpt.restore(tf2_ckpt_path).assert_consumed().run_restore_ops()\n",
        "\n",
        "    # Removing `assert_consumed` is fine if you want to skip the validation.\n",
        "    # ckpt.restore(tf2_ckpt_path).run_restore_ops()\n",
        "\n",
        "    print(\"Restored [a, b, c]: \", sess.run([a, b, c]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "unDPmL-kldr2"
      },
      "source": [
        "在 Eager 模式下："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "79S0zMAnfzx7"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(0.)\n",
        "b = tf.Variable(0.)\n",
        "c = tf.Variable(0.)\n",
        "print(\"Initialized [a, b, c]: \", [a.numpy(), b.numpy(), c.numpy()])\n",
        "\n",
        "# The keys \"scoped\" and \"unscoped\" are no longer relevant, but are used to\n",
        "# maintain compatibility with the saved checkpoints.\n",
        "ckpt = tf.train.Checkpoint(unscoped=[a, b], scoped=[c])\n",
        "\n",
        "ckpt.restore(tf2_ckpt_path).assert_consumed().run_restore_ops()\n",
        "print(\"Restored [a, b, c]: \", [a.numpy(), b.numpy(), c.numpy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dKfNAr8l3aFg"
      },
      "source": [
        "## Estimator 中的 TF2 检查点\n",
        "\n",
        "上面的部分介绍了如何在迁移模型时保持检查点兼容性。这些概念也适用于 Estimator 模型，尽管保存/加载检查点的方式略有不同。当迁移 Estimator 模型以使用 TF2 API 时，您可能希望*在模型仍使用 Estimator 时*从 TF1 切换到 TF2 检查点。本部分介绍如何实现此目的。\n",
        "\n",
        "[`tf.estimator.Estimator`](https://tensorflow.google.cn/api_docs/python/tf/estimator/Estimator) 和 [`MonitoredSession`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/MonitoredSession) 有一种称为 `scaffold` 的保存机制，即 [`tf.compat.v1.train.Scaffold`](https://tensorflow.google.cn/api_docs/python/tf/compat/v1/train/Scaffold) 对象。`Scaffold` 可以包含 `tf1.train.Saver` 或 `tf.train.Checkpoint`，它使 `Estimator` 和 `MonitoredSession` 能够保存 TF1 或 TF2 样式的检查点。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D8AT_oO-5TXU"
      },
      "outputs": [],
      "source": [
        "# A model_fn that saves a TF1 checkpoint\n",
        "def model_fn_tf1_ckpt(features, labels, mode):\n",
        "  # This model adds 2 to the variable `v` in every train step.\n",
        "  train_step = tf1.train.get_or_create_global_step()\n",
        "  v = tf1.get_variable('var', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(0))\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode,\n",
        "      predictions=v,\n",
        "      train_op=tf.group(v.assign_add(2), train_step.assign_add(1)),\n",
        "      loss=tf.constant(1.),\n",
        "      scaffold=None\n",
        "  )\n",
        "\n",
        "!rm -rf est-tf1\n",
        "est = tf.estimator.Estimator(model_fn_tf1_ckpt, 'est-tf1')\n",
        "\n",
        "def train_fn():\n",
        "  return tf.data.Dataset.from_tensor_slices(([1,2,3], [4,5,6]))\n",
        "est.train(train_fn, steps=1)\n",
        "\n",
        "latest_checkpoint = tf.train.latest_checkpoint('est-tf1')\n",
        "print_checkpoint(latest_checkpoint)  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ttH6cUrl7jK2"
      },
      "outputs": [],
      "source": [
        "# A model_fn that saves a TF2 checkpoint\n",
        "def model_fn_tf2_ckpt(features, labels, mode):\n",
        "  # This model adds 2 to the variable `v` in every train step.\n",
        "  train_step = tf1.train.get_or_create_global_step()\n",
        "  v = tf1.get_variable('var', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(0))\n",
        "  ckpt = tf.train.Checkpoint(var_list={'var': v}, step=train_step)\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode,\n",
        "      predictions=v,\n",
        "      train_op=tf.group(v.assign_add(2), train_step.assign_add(1)),\n",
        "      loss=tf.constant(1.),\n",
        "      scaffold=tf1.train.Scaffold(saver=ckpt)\n",
        "  )\n",
        "\n",
        "!rm -rf est-tf2\n",
        "est = tf.estimator.Estimator(model_fn_tf2_ckpt, 'est-tf2',\n",
        "                             warm_start_from='est-tf1')\n",
        "\n",
        "def train_fn():\n",
        "  return tf.data.Dataset.from_tensor_slices(([1,2,3], [4,5,6]))\n",
        "est.train(train_fn, steps=1)\n",
        "\n",
        "latest_checkpoint = tf.train.latest_checkpoint('est-tf2')\n",
        "print_checkpoint(latest_checkpoint)  \n",
        "\n",
        "assert est.get_variable_value('var_list/var/.ATTRIBUTES/VARIABLE_VALUE') == 4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hYVYgahE8daL"
      },
      "source": [
        "在从 `est-tf1` 热启动之后，`v` 的最终值应当是 `16`，随后再训练 5 步。训练步的值不会从 `warm_start` 检查点结转。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pq8EjblQUIA2"
      },
      "source": [
        "## 检查点 Keras\n",
        "\n",
        "使用 Keras 构建的模型仍然使用 `tf1.train.Saver` 和 `tf.train.Checkpoint` 来加载既有的权重。当您的模型完全迁移后，请切换为使用 `model.save_weights` 和 `model.load_weights`，尤其是当您在训练时使用 `ModelCheckpoint` 回调时。\n",
        "\n",
        "关于检查点和 Keras，您需要了解以下信息：\n",
        "\n",
        "**初始化与构建**\n",
        "\n",
        "Keras 模型和层必须经过**两个步骤**才能完全创建。首先是 Python 对象的*初始化*：`layer = tf.keras.layers.Dense(x)`。其次是*构建*步骤，此过程实际会创建大部分权重：`layer.build(input_shape)`。此外，您还可以通过调用或运行单个 `train`、`eval` 或 `predict` 步骤（仅限第一次）来构建模型。\n",
        "\n",
        "如果您发现 `model.load_weights(path).assert_consumed()` 引发错误，则很可能是模型/层尚未构建。\n",
        "\n",
        "**Keras 使用 TF2 检查点**\n",
        "\n",
        "`tf.train.Checkpoint(model).write` 等效于 `model.save_weights`。与 `tf.train.Checkpoint(model).read` 和 `model.load_weights` 相同。请注意，`Checkpoint(model) != Checkpoint(model=model)`。\n",
        "\n",
        "**TF2 检查点可与 Keras 的 `build()` 步骤一起使用**\n",
        "\n",
        "`tf.train.Checkpoint.restore` 有一种称为*延迟恢复*的机制，它允许在尚未创建变量时使用 `tf.Module` 和 Keras 对象存储变量值。这允许*已初始化的*模型加载权重并在之后*构建*。\n",
        "\n",
        "```\n",
        "m = YourKerasModel()\n",
        "status = m.load_weights(path)\n",
        "\n",
        "# This call builds the model. The variables are created with the restored\n",
        "# values.\n",
        "m.predict(inputs)\n",
        "\n",
        "status.assert_consumed()\n",
        "```\n",
        "\n",
        "由于存在这种机制，我们强烈建议您将 TF2 检查点加载 API 与 Keras 模型一起使用（即使在将既有的 TF1 检查点恢复到[模型映射填充码](./model_mapping.ipynb)中时）。有关详情，请参阅[检查点指南](https://tensorflow.google.cn/guide/checkpoint#delayed_restorations)。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xO2NucRtqMm6"
      },
      "source": [
        "## 代码段\n",
        "\n",
        "下面的代码段显示了检查点保存 API 中的 TF1/TF2 版本兼容性。 "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C3SSc74olkX3"
      },
      "source": [
        "### 在 TF2 中保存 TF1 检查点\n",
        "\n",
        "<a name=\"save-tf1-in-tf2\"></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t2ZPk8BPloE1"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(1.0, name='a')\n",
        "b = tf.Variable(2.0, name='b')\n",
        "with tf.name_scope('scoped'):\n",
        "  c = tf.Variable(3.0, name='c')\n",
        "\n",
        "saver = tf1.train.Saver(var_list=[a, b, c])\n",
        "path = saver.save(sess=None, save_path='tf1-ckpt-saved-in-eager')\n",
        "print_checkpoint(path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BxyN5khVjhmA"
      },
      "source": [
        "### 在 TF2 中加载 TF1 检查点\n",
        "\n",
        "<a name=\"load-tf1-in-tf2\"></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z5kSXy3FmA79"
      },
      "outputs": [],
      "source": [
        "a = tf.Variable(0., name='a')\n",
        "b = tf.Variable(0., name='b')\n",
        "with tf.name_scope('scoped'):\n",
        "  c = tf.Variable(0., name='c')\n",
        "print(\"Initialized [a, b, c]: \", [a.numpy(), b.numpy(), c.numpy()])\n",
        "saver = tf1.train.Saver(var_list=[a, b, c])\n",
        "saver.restore(sess=None, save_path='tf1-ckpt-saved-in-eager')\n",
        "print(\"Restored [a, b, c]: \", [a.numpy(), b.numpy(), c.numpy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ul3V4pEwloeN"
      },
      "source": [
        "### 在 TF1 中保存 TF2 检查点"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UhuP_2EIlRm4"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(1))\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(2))\n",
        "  with tf1.variable_scope('scoped'):\n",
        "    c = tf1.get_variable('c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.constant_initializer(3))\n",
        "  with tf1.Session() as sess:\n",
        "    sess.run(tf1.global_variables_initializer())\n",
        "    ckpt = tf.train.Checkpoint(\n",
        "        var_list={v.name.split(':')[0]: v for v in tf1.global_variables()})\n",
        "    tf2_in_tf1_path = ckpt.save('tf2-ckpt-saved-in-session')\n",
        "    print_checkpoint(tf2_in_tf1_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GiViCjCDgxhz"
      },
      "source": [
        "### 在 TF1 中加载 TF2 检查点\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j-4hIPZvmXlb"
      },
      "outputs": [],
      "source": [
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(0))\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(0))\n",
        "  with tf1.variable_scope('scoped'):\n",
        "    c = tf1.get_variable('c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.constant_initializer(0))\n",
        "  with tf1.Session() as sess:\n",
        "    sess.run(tf1.global_variables_initializer())\n",
        "    print(\"Initialized [a, b, c]: \", sess.run([a, b, c]))\n",
        "    ckpt = tf.train.Checkpoint(\n",
        "        var_list={v.name.split(':')[0]: v for v in tf1.global_variables()})\n",
        "    ckpt.restore('tf2-ckpt-saved-in-session-1').run_restore_ops()\n",
        "    print(\"Restored [a, b, c]: \", sess.run([a, b, c]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oRrSE2X6sgAM"
      },
      "source": [
        "## 检查点转换\n",
        "\n",
        "<a name=\"checkpoint-conversion\"></a>\n",
        "\n",
        "可以通过加载和重新保存检查点在 TF1 和 TF2 之间转换检查点。另一种方式是 `tf.train.load_checkpoint`，如下面的代码所示。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9KByaLous4q"
      },
      "source": [
        "### 将 TF1 检查点转换为 TF2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NG8grCv2smAb"
      },
      "outputs": [],
      "source": [
        "def convert_tf1_to_tf2(checkpoint_path, output_prefix):\n",
        "  \"\"\"Converts a TF1 checkpoint to TF2.\n",
        "\n",
        "  To load the converted checkpoint, you must build a dictionary that maps\n",
        "  variable names to variable objects.\n",
        "  ```\n",
        "  ckpt = tf.train.Checkpoint(vars={name: variable})  \n",
        "  ckpt.restore(converted_ckpt_path)\n",
        "  ```\n",
        "\n",
        "  Args:\n",
        "    checkpoint_path: Path to the TF1 checkpoint.\n",
        "    output_prefix: Path prefix to the converted checkpoint.\n",
        "\n",
        "  Returns:\n",
        "    Path to the converted checkpoint.\n",
        "  \"\"\"\n",
        "  vars = {}\n",
        "  reader = tf.train.load_checkpoint(checkpoint_path)\n",
        "  dtypes = reader.get_variable_to_dtype_map()\n",
        "  for key in dtypes.keys():\n",
        "    vars[key] = tf.Variable(reader.get_tensor(key))\n",
        "  return tf.train.Checkpoint(vars=vars).save(output_prefix)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TyvqK6Sb3dad"
      },
      "source": [
        "转换代码段 `Save a TF1 checkpoint in TF2` 中保存的检查点："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gcHLN4lPvYvw"
      },
      "outputs": [],
      "source": [
        "# Make sure to run the snippet in `Save a TF1 checkpoint in TF2`.\n",
        "print_checkpoint('tf1-ckpt-saved-in-eager')\n",
        "converted_path = convert_tf1_to_tf2('tf1-ckpt-saved-in-eager', \n",
        "                                     'converted-tf1-to-tf2')\n",
        "print(\"\\n[Converted]\")\n",
        "print_checkpoint(converted_path)\n",
        "\n",
        "# Try loading the converted checkpoint.\n",
        "a = tf.Variable(0.)\n",
        "b = tf.Variable(0.)\n",
        "c = tf.Variable(0.)\n",
        "ckpt = tf.train.Checkpoint(vars={'a': a, 'b': b, 'scoped/c': c})\n",
        "ckpt.restore(converted_path).assert_consumed()\n",
        "print(\"\\nRestored [a, b, c]: \", [a.numpy(), b.numpy(), c.numpy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fokg6ybZvE20"
      },
      "source": [
        "### 将 TF2 检查点转换为 TF1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NPQsXQveuQiC"
      },
      "outputs": [],
      "source": [
        "def convert_tf2_to_tf1(checkpoint_path, output_prefix):\n",
        "  \"\"\"Converts a TF2 checkpoint to TF1.\n",
        "\n",
        "  The checkpoint must be saved using a \n",
        "  `tf.train.Checkpoint(var_list={name: variable})`\n",
        "\n",
        "  To load the converted checkpoint with `tf.compat.v1.Saver`:\n",
        "  ```\n",
        "  saver = tf.compat.v1.train.Saver(var_list={name: variable}) \n",
        "\n",
        "  # An alternative, if the variable names match the keys:\n",
        "  saver = tf.compat.v1.train.Saver(var_list=[variables]) \n",
        "  saver.restore(sess, output_path)\n",
        "  ```\n",
        "  \"\"\"\n",
        "  vars = {}\n",
        "  reader = tf.train.load_checkpoint(checkpoint_path)\n",
        "  dtypes = reader.get_variable_to_dtype_map()\n",
        "  for key in dtypes.keys():\n",
        "    # Get the \"name\" from the \n",
        "    if key.startswith('var_list/'):\n",
        "      var_name = key.split('/')[1]\n",
        "      # TF2 checkpoint keys use '/', so if they appear in the user-defined name,\n",
        "      # they are escaped to '.S'.\n",
        "      var_name = var_name.replace('.S', '/')\n",
        "      vars[var_name] = tf.Variable(reader.get_tensor(key))\n",
        "  \n",
        "  return tf1.train.Saver(var_list=vars).save(sess=None, save_path=output_prefix)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VjZD_OSf1mKX"
      },
      "source": [
        "转换代码段 `Save a TF2 checkpoint in TF1` 中保存的检查点："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vc1MVeV6z2DB"
      },
      "outputs": [],
      "source": [
        "# Make sure to run the snippet in `Save a TF2 checkpoint in TF1`.\n",
        "print_checkpoint('tf2-ckpt-saved-in-session-1')\n",
        "converted_path = convert_tf2_to_tf1('tf2-ckpt-saved-in-session-1',\n",
        "                                    'converted-tf2-to-tf1')\n",
        "print(\"\\n[Converted]\")\n",
        "print_checkpoint(converted_path)\n",
        "\n",
        "# Try loading the converted checkpoint.\n",
        "with tf.Graph().as_default() as g:\n",
        "  a = tf1.get_variable('a', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(0))\n",
        "  b = tf1.get_variable('b', shape=[], dtype=tf.float32, \n",
        "                       initializer=tf1.constant_initializer(0))\n",
        "  with tf1.variable_scope('scoped'):\n",
        "    c = tf1.get_variable('c', shape=[], dtype=tf.float32, \n",
        "                        initializer=tf1.constant_initializer(0))\n",
        "  with tf1.Session() as sess:\n",
        "    saver = tf1.train.Saver([a, b, c])\n",
        "    saver.restore(sess, converted_path)\n",
        "    print(\"\\nRestored [a, b, c]: \", sess.run([a, b, c]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JBMfArLQ0jb-"
      },
      "source": [
        "## 相关指南\n",
        "\n",
        "- [验证数值等价关系和正确性](./validate_correctness.ipynb)\n",
        "- [模型映射指南](./model_mapping.ipynb)和 `tf.compat.v1.keras.utils.track_tf1_style_variables`\n",
        "- [TF2 检查点指南](https://tensorflow.google.cn/guide/checkpoint)。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "migrating_checkpoints.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}