在 TF2 工作流程中使用 TF1.x 模型

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

本指南概述并举例说明了您可以用来在 TF2 工作流程中使用现有 TF1.x 模型的 建模代码垫片，例如急切执行、tf.function 和分布式策略，只需对您的建模代码进行最小的更改。

使用范围

本指南中描述的垫片专为依赖以下内容的 TF1.x 模型而设计

1. tf.compat.v1.get_variable 和 tf.compat.v1.variable_scope 来控制变量创建和重用，以及
2. 基于图集合的 API，例如 tf.compat.v1.global_variables()、tf.compat.v1.trainable_variables、tf.compat.v1.losses.get_regularization_losses() 和 tf.compat.v1.get_collection() 来跟踪权重和正则化损失

这包括大多数在 tf.compat.v1.layer、tf.contrib.layers API 和 TensorFlow-Slim 之上构建的模型。

对于以下 TF1.x 模型，不需要 使用垫片

1. 独立的 Keras 模型，这些模型已经通过 model.trainable_weights 和 model.losses 分别跟踪了所有可训练权重和正则化损失。
2. tf.Module，这些模型已经通过 module.trainable_variables 跟踪了所有可训练权重，并且仅在权重尚未创建时才创建权重。

这些模型很可能在 TF2 中使用急切执行和 tf.function 能够开箱即用。

设置

导入 TensorFlow 和其他依赖项。

pip uninstall -y -q tensorflow

# Install tf-nightly as the DeterministicRandomTestTool is available only in
# Tensorflow 2.8

pip install -q tf-nightly

import tensorflow as tf
import tensorflow.compat.v1 as v1
import sys
import numpy as np

from contextlib import contextmanager

track_tf1_style_variables 装饰器

本指南中描述的关键垫片是 tf.compat.v1.keras.utils.track_tf1_style_variables，这是一个装饰器，您可以在属于 tf.keras.layers.Layer 和 tf.Module 的方法中使用它来跟踪 TF1.x 风格的权重并捕获正则化损失。

使用 tf.compat.v1.keras.utils.track_tf1_style_variables 装饰 tf.keras.layers.Layer 或 tf.Module 的调用方法，允许通过 tf.compat.v1.get_variable（以及扩展的 tf.compat.v1.layers）在装饰的方法内部正确地进行变量创建和重用，而不是在每次调用时都创建一个新变量。它还会导致层或模块隐式跟踪通过 get_variable 在装饰的方法内部创建或访问的任何权重。

除了在标准 layer.variable/module.variable/等属性下跟踪权重本身之外，如果该方法属于 tf.keras.layers.Layer，那么通过 get_variable 或 tf.compat.v1.layers 的 regularizer 参数指定的任何正则化损失都将由层在标准 layer.losses 属性下跟踪。

这种跟踪机制使您能够在 TF2 中使用 Keras 层或 tf.Module 中使用大量 TF1.x 风格的模型正向传递代码，即使启用了 TF2 行为。

使用示例

以下使用示例演示了用于装饰 tf.keras.layers.Layer 方法的建模垫片，但除了它们与 Keras 功能的特定交互之外，它们也适用于装饰 tf.Module 方法。

使用 tf.compat.v1.get_variable 构建的层

假设您有一个直接在 tf.compat.v1.get_variable 之上实现的层，如下所示

def dense(self, inputs, units):
  out = inputs
  with tf.compat.v1.variable_scope("dense"):
    # The weights are created with a `regularizer`,
    kernel = tf.compat.v1.get_variable(
        shape=[out.shape[-1], units],
        regularizer=tf.keras.regularizers.L2(),
        initializer=tf.compat.v1.initializers.glorot_normal,
        name="kernel")
    bias = tf.compat.v1.get_variable(
        shape=[units,],
        initializer=tf.compat.v1.initializers.zeros,
        name="bias")
    out = tf.linalg.matmul(out, kernel)
    out = tf.compat.v1.nn.bias_add(out, bias)
  return out

使用垫片将其转换为层，并在输入上调用它。

class DenseLayer(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs):
    out = inputs
    with tf.compat.v1.variable_scope("dense"):
      # The weights are created with a `regularizer`,
      # so the layer should track their regularization losses
      kernel = tf.compat.v1.get_variable(
          shape=[out.shape[-1], self.units],
          regularizer=tf.keras.regularizers.L2(),
          initializer=tf.compat.v1.initializers.glorot_normal,
          name="kernel")
      bias = tf.compat.v1.get_variable(
          shape=[self.units,],
          initializer=tf.compat.v1.initializers.zeros,
          name="bias")
      out = tf.linalg.matmul(out, kernel)
      out = tf.compat.v1.nn.bias_add(out, bias)
    return out

layer = DenseLayer(10)
x = tf.random.normal(shape=(8, 20))
layer(x)

像标准 Keras 层一样访问跟踪的变量和捕获的正则化损失。

layer.trainable_variables
layer.losses

要查看权重在每次调用层时是否被重用，请将所有权重设置为零，然后再次调用该层。

print("Resetting variables to zero:", [var.name for var in layer.trainable_variables])

for var in layer.trainable_variables:
  var.assign(var * 0.0)

# Note: layer.losses is not a live view and
# will get reset only at each layer call
print("layer.losses:", layer.losses)
print("calling layer again.")
out = layer(x)
print("layer.losses: ", layer.losses)
out

您也可以在 Keras 函数模型构建中直接使用转换后的层。

inputs = tf.keras.Input(shape=(20))
outputs = DenseLayer(10)(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

x = tf.random.normal(shape=(8, 20))
model(x)

# Access the model variables and regularization losses
model.weights
model.losses

使用 tf.compat.v1.layers 构建的模型

假设您有一个直接在 tf.compat.v1.layers 之上实现的层或模型，如下所示

def model(self, inputs, units):
  with tf.compat.v1.variable_scope('model'):
    out = tf.compat.v1.layers.conv2d(
        inputs, 3, 3,
        kernel_regularizer="l2")
    out = tf.compat.v1.layers.flatten(out)
    out = tf.compat.v1.layers.dense(
        out, units,
        kernel_regularizer="l2")
    return out

使用垫片将其转换为层，并在输入上调用它。

class CompatV1LayerModel(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs):
    with tf.compat.v1.variable_scope('model'):
      out = tf.compat.v1.layers.conv2d(
          inputs, 3, 3,
          kernel_regularizer="l2")
      out = tf.compat.v1.layers.flatten(out)
      out = tf.compat.v1.layers.dense(
          out, self.units,
          kernel_regularizer="l2")
      return out

layer = CompatV1LayerModel(10)
x = tf.random.normal(shape=(8, 5, 5, 5))
layer(x)

像标准 Keras 层一样访问跟踪的变量和捕获的正则化损失。

layer.trainable_variables
layer.losses

要查看权重在每次调用层时是否被重用，请将所有权重设置为零，然后再次调用该层。

print("Resetting variables to zero:", [var.name for var in layer.trainable_variables])

for var in layer.trainable_variables:
  var.assign(var * 0.0)

out = layer(x)
print("layer.losses: ", layer.losses)
out

您也可以在 Keras 函数模型构建中直接使用转换后的层。

inputs = tf.keras.Input(shape=(5, 5, 5))
outputs = CompatV1LayerModel(10)(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

x = tf.random.normal(shape=(8, 5, 5, 5))
model(x)

# Access the model variables and regularization losses
model.weights
model.losses

捕获批次归一化更新和模型 training 参数

在 TF1.x 中，您像这样执行批次归一化

  x_norm = tf.compat.v1.layers.batch_normalization(x, training=training)

  # ...

  update_ops = tf.compat.v1.get_collection(tf.GraphKeys.UPDATE_OPS)
  train_op = optimizer.minimize(loss)
  train_op = tf.group([train_op, update_ops])

请注意

1. 批次归一化移动平均更新由 get_collection 跟踪，该函数是在与层分离的情况下调用的
2. tf.compat.v1.layers.batch_normalization 需要一个 training 参数（通常在使用 TF-Slim 批次归一化层时称为 is_training）

在 TF2 中，由于 急切执行 和自动控制依赖项，批次归一化移动平均更新将立即执行。无需从更新集合中单独收集它们，并将它们添加为显式控制依赖项。

此外，如果您为 tf.keras.layers.Layer 的前向传递方法提供一个 training 参数，Keras 将能够像对任何其他层一样将当前训练阶段和任何嵌套层传递给它。有关 Keras 如何处理 training 参数的更多信息，请参阅 tf.keras.Model 的 API 文档。

如果您正在装饰 tf.Module 方法，您需要确保手动传递所有 training 参数（如果需要）。但是，批次归一化移动平均更新仍将自动应用，无需显式控制依赖项。

以下代码片段演示了如何在垫片中嵌入批次归一化层，以及如何在 Keras 模型中使用它（适用于 tf.keras.layers.Layer）。

class CompatV1BatchNorm(tf.keras.layers.Layer):

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs, training=None):
    print("Forward pass called with `training` =", training)
    with v1.variable_scope('batch_norm_layer'):
      return v1.layers.batch_normalization(x, training=training)

print("Constructing model")
inputs = tf.keras.Input(shape=(5, 5, 5))
outputs = CompatV1BatchNorm()(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

print("Calling model in inference mode")
x = tf.random.normal(shape=(8, 5, 5, 5))
model(x, training=False)

print("Moving average variables before training: ",
      {var.name: var.read_value() for var in model.non_trainable_variables})

# Notice that when running TF2 and eager execution, the batchnorm layer directly
# updates the moving averages while training without needing any extra control
# dependencies
print("calling model in training mode")
model(x, training=True)

print("Moving average variables after training: ",
      {var.name: var.read_value() for var in model.non_trainable_variables})

基于变量作用域的变量重用

前向传递中基于 get_variable 的任何变量创建都将保持与 TF1.x 中变量作用域相同的变量命名和重用语义。只要您为任何具有自动生成名称的 tf.compat.v1.layers 至少有一个非空外部作用域，如上所述，就会如此。

急切执行和 tf.function

如上所示，tf.keras.layers.Layer 和 tf.Module 的装饰方法在急切执行中运行，并且也与 tf.function 兼容。这意味着您可以使用 pdb 和其他交互式工具来逐步执行您的前向传递，因为它正在运行。

分布策略

在 @track_tf1_style_variables 装饰的层或模块方法内部对 get_variable 的调用在幕后使用标准 tf.Variable 变量创建。这意味着您可以将它们与 tf.distribute 提供的各种分布策略一起使用，例如 MirroredStrategy 和 TPUStrategy。

在装饰的调用中嵌套 tf.Variable、tf.Module、tf.keras.layers 和 tf.keras.models

在 tf.compat.v1.keras.utils.track_tf1_style_variables 中装饰您的层调用只会添加通过 tf.compat.v1.get_variable 创建（和重用）的变量的自动隐式跟踪。它不会直接捕获由 tf.Variable 调用创建的权重，例如那些由典型 Keras 层和大多数 tf.Module 使用的权重。本节介绍如何处理这些嵌套情况。

（预先存在的用法）tf.keras.layers 和 tf.keras.models

对于嵌套 Keras 层和模型的预先存在的用法，请使用 tf.compat.v1.keras.utils.get_or_create_layer。这仅推荐用于简化现有 TF1.x 嵌套 Keras 用法的迁移；新代码应使用下面描述的显式属性设置来设置 tf.Variables 和 tf.Modules。

要使用 tf.compat.v1.keras.utils.get_or_create_layer，请将构建嵌套模型的代码包装到一个方法中，并将其传递给该方法。示例

class NestedModel(tf.keras.Model):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units

  def build_model(self):
    inp = tf.keras.Input(shape=(5, 5))
    dense_layer = tf.keras.layers.Dense(
        10, name="dense", kernel_regularizer="l2",
        kernel_initializer=tf.compat.v1.ones_initializer())
    model = tf.keras.Model(inputs=inp, outputs=dense_layer(inp))
    return model

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs):
    # Get or create a nested model without assigning it as an explicit property
    model = tf.compat.v1.keras.utils.get_or_create_layer(
        "dense_model", self.build_model)
    return model(inputs)

layer = NestedModel(10)
layer(tf.ones(shape=(5,5)))

此方法确保这些嵌套层被 tensorflow 正确重用和跟踪。请注意，@track_tf1_style_variables 装饰器仍然需要在适当的方法上。传递给 get_or_create_layer 的模型构建器方法（在本例中为 self.build_model）不应接受任何参数。

权重被跟踪

assert len(layer.weights) == 2
weights = {x.name: x for x in layer.variables}

assert set(weights.keys()) == {"dense/bias:0", "dense/kernel:0"}

layer.weights

以及正则化损失

tf.add_n(layer.losses)

增量迁移：tf.Variables 和 tf.Modules

如果您需要在装饰的方法中嵌入 tf.Variable 调用或 tf.Module（例如，如果您正在遵循本指南后面介绍的增量迁移到非遗留 TF2 API），您仍然需要显式跟踪这些，并满足以下要求

• 显式确保变量/模块/层仅创建一次
• 像定义 典型模块或层 一样，将它们显式地附加为实例属性
• 在后续调用中显式重用已创建的对象

这确保了权重不会在每次调用时都重新创建，并且被正确重用。此外，这也确保了现有权重和正则化损失被跟踪。

以下是如何执行此操作的示例

class NestedLayer(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def __call__(self, inputs):
    out = inputs
    with tf.compat.v1.variable_scope("inner_dense"):
      # The weights are created with a `regularizer`,
      # so the layer should track their regularization losses
      kernel = tf.compat.v1.get_variable(
          shape=[out.shape[-1], self.units],
          regularizer=tf.keras.regularizers.L2(),
          initializer=tf.compat.v1.initializers.glorot_normal,
          name="kernel")
      bias = tf.compat.v1.get_variable(
          shape=[self.units,],
          initializer=tf.compat.v1.initializers.zeros,
          name="bias")
      out = tf.linalg.matmul(out, kernel)
      out = tf.compat.v1.nn.bias_add(out, bias)
    return out

class WrappedDenseLayer(tf.keras.layers.Layer):

  def __init__(self, units, **kwargs):
    super().__init__(**kwargs)
    self.units = units
    # Only create the nested tf.variable/module/layer/model
    # once, and then reuse it each time!
    self._dense_layer = NestedLayer(self.units)

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs):
    with tf.compat.v1.variable_scope('outer'):
      outputs = tf.compat.v1.layers.dense(inputs, 3)
      outputs = tf.compat.v1.layers.dense(inputs, 4)
      return self._dense_layer(outputs)

layer = WrappedDenseLayer(10)

layer(tf.ones(shape=(5, 5)))

请注意，即使嵌套模块用 track_tf1_style_variables 装饰器装饰，也需要显式跟踪它。这是因为每个具有装饰方法的模块/层都有与其关联的自己的变量存储。

权重被正确跟踪

assert len(layer.weights) == 6
weights = {x.name: x for x in layer.variables}

assert set(weights.keys()) == {"outer/inner_dense/bias:0",
                               "outer/inner_dense/kernel:0",
                               "outer/dense/bias:0",
                               "outer/dense/kernel:0",
                               "outer/dense_1/bias:0",
                               "outer/dense_1/kernel:0"}

layer.trainable_weights

以及正则化损失

layer.losses

请注意，如果 NestedLayer 是一个非 Keras tf.Module，则变量仍将被跟踪，但正则化损失不会自动跟踪，因此您必须单独显式跟踪它们。

有关变量名称的指南

显式 tf.Variable 调用和 Keras 层使用与您可能习惯的 get_variable 和 variable_scopes 组合不同的层名称/变量名称自动生成机制。虽然垫片将使您的变量名称在从 TF1.x 图到 TF2 急切执行和 tf.function 时匹配，但它无法保证您在方法装饰器中嵌入的 tf.Variable 调用和 Keras 层生成的变量名称相同。在 TF2 急切执行和 tf.function 中，多个变量甚至可能共享相同的名称。

在遵循本指南后面介绍的有关验证正确性和映射 TF1.x 检查点的部分时，您应该特别注意这一点。

在装饰的方法中使用 tf.compat.v1.make_template

强烈建议您直接使用 tf.compat.v1.keras.utils.track_tf1_style_variables，而不是使用 tf.compat.v1.make_template，因为它是在 TF2 之上的更薄的一层.

对于已经依赖 tf.compat.v1.make_template 的先前 TF1.x 代码，请遵循本节中的指南。

因为 tf.compat.v1.make_template 包装了使用 get_variable 的代码，所以 track_tf1_style_variables 装饰器允许您在层调用中使用这些模板，并成功跟踪权重和正则化损失。

但是，请确保只调用 make_template 一次，然后在每次层调用中重用相同的模板。否则，每次调用层时都会创建一个新的模板，以及一组新的变量。

例如，

class CompatV1TemplateScaleByY(tf.keras.layers.Layer):

  def __init__(self, **kwargs):
    super().__init__(**kwargs)
    def my_op(x, scalar_name):
      var1 = tf.compat.v1.get_variable(scalar_name,
                            shape=[],
                            regularizer=tf.compat.v1.keras.regularizers.L2(),
                            initializer=tf.compat.v1.constant_initializer(1.5))
      return x * var1
    self.scale_by_y = tf.compat.v1.make_template('scale_by_y', my_op, scalar_name='y')

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs):
    with tf.compat.v1.variable_scope('layer'):
      # Using a scope ensures the `scale_by_y` name will not be incremented
      # for each instantiation of the layer.
      return self.scale_by_y(inputs)

layer = CompatV1TemplateScaleByY()

out = layer(tf.ones(shape=(2, 3)))
print("weights:", layer.weights)
print("regularization loss:", layer.losses)
print("output:", out)

增量迁移到原生 TF2

如前所述，track_tf1_style_variables 允许您在同一个装饰模块/层中混合使用 TF2 风格的面向对象 tf.Variable/tf.keras.layers.Layer/tf.Module 使用方法和传统 tf.compat.v1.get_variable/tf.compat.v1.layers 风格的使用方法。

这意味着，在您将 TF1.x 模型完全转换为 TF2 兼容后，您可以使用原生（非 tf.compat.v1）TF2 API 编写所有新的模型组件，并使其与您的旧代码互操作。

但是，如果您继续修改旧的模型组件，您也可以选择将传统的 tf.compat.v1 使用方法逐步切换到纯粹的原生面向对象 API，这些 API 是为新编写的 TF2 代码推荐的。

tf.compat.v1.get_variable 的使用可以替换为 self.add_weight 调用（如果您正在装饰 Keras 层/模型），或者替换为 tf.Variable 调用（如果您正在装饰 Keras 对象或 tf.Module）。

功能风格和面向对象的 tf.compat.v1.layers 通常可以替换为等效的 tf.keras.layers 层，无需更改任何参数。

在您逐步迁移到纯粹的原生 API 的过程中，您也可以考虑将模型的某些部分或常见模式分成单独的层/模块，这些层/模块本身可能使用 track_tf1_style_variables。

关于 Slim 和 contrib.layers 的说明

大量的旧 TF 1.x 代码使用 Slim 库，该库在 TF 1.x 中打包为 tf.contrib.layers。将使用 Slim 的代码转换为原生 TF 2 比将 v1.layers 转换为原生 TF 2 更复杂。实际上，您可能需要先将 Slim 代码转换为 v1.layers，然后再转换为 Keras。以下是一些关于转换 Slim 代码的一般指南。

• 确保所有参数都是显式的。如果可能，删除 arg_scopes。如果您仍然需要使用它们，请将 normalizer_fn 和 activation_fn 分成各自的层。
• 可分离卷积层映射到一个或多个不同的 Keras 层（深度方向、逐点方向和可分离 Keras 层）。
• Slim 和 v1.layers 具有不同的参数名称和默认值。
• 请注意，某些参数具有不同的比例。

忽略检查点兼容性的原生 TF2 迁移

以下代码示例演示了将模型逐步迁移到纯粹的原生 API，而无需考虑检查点兼容性。

class CompatModel(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs, training=None):
    with tf.compat.v1.variable_scope('model'):
      out = tf.compat.v1.layers.conv2d(
          inputs, 3, 3,
          kernel_regularizer="l2")
      out = tf.compat.v1.layers.flatten(out)
      out = tf.compat.v1.layers.dropout(out, training=training)
      out = tf.compat.v1.layers.dense(
          out, self.units,
          kernel_regularizer="l2")
      return out

接下来，以分段的方式将 compat.v1 API 替换为其原生面向对象等效项。首先将卷积层切换到在层构造函数中创建的 Keras 对象。

class PartiallyMigratedModel(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units
    self.conv_layer = tf.keras.layers.Conv2D(
      3, 3,
      kernel_regularizer="l2")

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs, training=None):
    with tf.compat.v1.variable_scope('model'):
      out = self.conv_layer(inputs)
      out = tf.compat.v1.layers.flatten(out)
      out = tf.compat.v1.layers.dropout(out, training=training)
      out = tf.compat.v1.layers.dense(
          out, self.units,
          kernel_regularizer="l2")
      return out

使用 v1.keras.utils.DeterministicRandomTestTool 类来验证此增量更改是否使模型的行为与之前相同。

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
with random_tool.scope():
  tf.keras.utils.set_random_seed(42)
  layer = CompatModel(10)

  inputs = tf.random.normal(shape=(10, 5, 5, 5))
  original_output = layer(inputs)

  # Grab the regularization loss as well
  original_regularization_loss = tf.math.add_n(layer.losses)

print(original_regularization_loss)

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
with random_tool.scope():
  tf.keras.utils.set_random_seed(42)
  layer = PartiallyMigratedModel(10)

  inputs = tf.random.normal(shape=(10, 5, 5, 5))
  migrated_output = layer(inputs)

  # Grab the regularization loss as well
  migrated_regularization_loss = tf.math.add_n(layer.losses)

print(migrated_regularization_loss)

# Verify that the regularization loss and output both match
np.testing.assert_allclose(original_regularization_loss.numpy(), migrated_regularization_loss.numpy())
np.testing.assert_allclose(original_output.numpy(), migrated_output.numpy())

现在，您已经将所有单独的 compat.v1.layers 替换为原生 Keras 层。

class NearlyFullyNativeModel(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units
    self.conv_layer = tf.keras.layers.Conv2D(
      3, 3,
      kernel_regularizer="l2")
    self.flatten_layer = tf.keras.layers.Flatten()
    self.dense_layer = tf.keras.layers.Dense(
      self.units,
      kernel_regularizer="l2")

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs):
    with tf.compat.v1.variable_scope('model'):
      out = self.conv_layer(inputs)
      out = self.flatten_layer(out)
      out = self.dense_layer(out)
      return out

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
with random_tool.scope():
  tf.keras.utils.set_random_seed(42)
  layer = NearlyFullyNativeModel(10)

  inputs = tf.random.normal(shape=(10, 5, 5, 5))
  migrated_output = layer(inputs)

  # Grab the regularization loss as well
  migrated_regularization_loss = tf.math.add_n(layer.losses)

print(migrated_regularization_loss)

# Verify that the regularization loss and output both match
np.testing.assert_allclose(original_regularization_loss.numpy(), migrated_regularization_loss.numpy())
np.testing.assert_allclose(original_output.numpy(), migrated_output.numpy())

最后，删除任何剩余的（不再需要的）variable_scope 使用方法和 track_tf1_style_variables 装饰器本身。

现在，您只剩下一个完全使用原生 API 的模型版本。

class FullyNativeModel(tf.keras.layers.Layer):

  def __init__(self, units, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.units = units
    self.conv_layer = tf.keras.layers.Conv2D(
      3, 3,
      kernel_regularizer="l2")
    self.flatten_layer = tf.keras.layers.Flatten()
    self.dense_layer = tf.keras.layers.Dense(
      self.units,
      kernel_regularizer="l2")

  def call(self, inputs):
    out = self.conv_layer(inputs)
    out = self.flatten_layer(out)
    out = self.dense_layer(out)
    return out

random_tool = v1.keras.utils.DeterministicRandomTestTool(mode='num_random_ops')
with random_tool.scope():
  tf.keras.utils.set_random_seed(42)
  layer = FullyNativeModel(10)

  inputs = tf.random.normal(shape=(10, 5, 5, 5))
  migrated_output = layer(inputs)

  # Grab the regularization loss as well
  migrated_regularization_loss = tf.math.add_n(layer.losses)

print(migrated_regularization_loss)

# Verify that the regularization loss and output both match
np.testing.assert_allclose(original_regularization_loss.numpy(), migrated_regularization_loss.numpy())
np.testing.assert_allclose(original_output.numpy(), migrated_output.numpy())

在迁移到原生 TF2 期间维护检查点兼容性

上述迁移到原生 TF2 API 的过程更改了变量名称（因为 Keras API 会生成非常不同的权重名称）以及指向模型中不同权重的面向对象路径。这些更改的影响是，它们将破坏任何现有的 TF1 风格的基于名称的检查点或 TF2 风格的面向对象检查点。

但是，在某些情况下，您可能能够使用您原始的基于名称的检查点，并使用类似于 重用 TF1.x 检查点指南 中详细介绍的方法来找到变量与其新名称的映射。

以下是一些使这成为可能的技巧

• 变量仍然都具有可以设置的 name 参数。
• Keras 模型也接受 name 参数，它们将此参数设置为其变量的前缀。
• v1.name_scope 函数可用于设置变量名称前缀。这与 tf.variable_scope 非常不同。它只影响名称，不跟踪变量和重用。

考虑到上述要点，以下代码示例演示了您可以调整到您的代码的工作流程，以逐步更新模型的一部分，同时更新检查点。

1. 首先将功能风格的 tf.compat.v1.layers 切换到其面向对象版本。

class FunctionalStyleCompatModel(tf.keras.layers.Layer):

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs, training=None):
    with tf.compat.v1.variable_scope('model'):
      out = tf.compat.v1.layers.conv2d(
          inputs, 3, 3,
          kernel_regularizer="l2")
      out = tf.compat.v1.layers.conv2d(
          out, 4, 4,
          kernel_regularizer="l2")
      out = tf.compat.v1.layers.conv2d(
          out, 5, 5,
          kernel_regularizer="l2")
      return out

layer = FunctionalStyleCompatModel()
layer(tf.ones(shape=(10, 10, 10, 10)))
[v.name for v in layer.weights]

1. 接下来，将 compat.v1.layer 对象和由 compat.v1.get_variable 创建的任何变量分配为 tf.keras.layers.Layer/tf.Module 对象的属性，该对象的 method 用 track_tf1_style_variables 装饰（请注意，任何面向对象的 TF2 风格的检查点现在将保存变量名称的路径和新的面向对象路径）。

class OOStyleCompatModel(tf.keras.layers.Layer):

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.conv_1 = tf.compat.v1.layers.Conv2D(
          3, 3,
          kernel_regularizer="l2")
    self.conv_2 = tf.compat.v1.layers.Conv2D(
          4, 4,
          kernel_regularizer="l2")

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs, training=None):
    with tf.compat.v1.variable_scope('model'):
      out = self.conv_1(inputs)
      out = self.conv_2(out)
      out = tf.compat.v1.layers.conv2d(
          out, 5, 5,
          kernel_regularizer="l2")
      return out

layer = OOStyleCompatModel()
layer(tf.ones(shape=(10, 10, 10, 10)))
[v.name for v in layer.weights]

1. 此时重新保存加载的检查点，以保存变量名称的路径（对于 compat.v1.layers）或面向对象对象图的路径。

weights = {v.name: v for v in layer.weights}
assert weights['model/conv2d/kernel:0'] is layer.conv_1.kernel
assert weights['model/conv2d_1/bias:0'] is layer.conv_2.bias

1. 现在，您可以将面向对象的 compat.v1.layers 替换为原生 Keras 层，同时仍然能够加载最近保存的检查点。确保您为剩余的 compat.v1.layers 保留变量名称，方法是仍然记录被替换层的自动生成的 variable_scopes。这些切换的层/变量现在将只使用检查点中变量的对象属性路径，而不是变量名称路径。

通常，您可以通过以下方式替换 compat.v1.get_variable 在附加到属性的变量中的使用方式

• 将它们切换到使用 tf.Variable，或者
• 使用 tf.keras.layers.Layer.add_weight 更新它们。请注意，如果您不是一次性切换所有层，这可能会更改为剩余的 compat.v1.layers（缺少 name 参数）自动生成的层/变量命名。如果是这种情况，您必须通过手动打开和关闭与已删除的 compat.v1.layer 生成的范围名称相对应的 variable_scope 来保持剩余的 compat.v1.layers 的变量名称相同。否则，现有检查点中的路径可能会冲突，检查点加载将表现不正常。

def record_scope(scope_name):
  """Record a variable_scope to make sure future ones get incremented."""
  with tf.compat.v1.variable_scope(scope_name):
    pass

class PartiallyNativeKerasLayersModel(tf.keras.layers.Layer):

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.conv_1 = tf.keras.layers.Conv2D(
          3, 3,
          kernel_regularizer="l2")
    self.conv_2 = tf.keras.layers.Conv2D(
          4, 4,
          kernel_regularizer="l2")

  @tf.compat.v1.keras.utils.track_tf1_style_variables
  def call(self, inputs, training=None):
    with tf.compat.v1.variable_scope('model'):
      out = self.conv_1(inputs)
      record_scope('conv2d') # Only needed if follow-on compat.v1.layers do not pass a `name` arg
      out = self.conv_2(out)
      record_scope('conv2d_1') # Only needed if follow-on compat.v1.layers do not pass a `name` arg
      out = tf.compat.v1.layers.conv2d(
          out, 5, 5,
          kernel_regularizer="l2")
      return out

layer = PartiallyNativeKerasLayersModel()
layer(tf.ones(shape=(10, 10, 10, 10)))
[v.name for v in layer.weights]

在构造变量后在此步骤中保存检查点将使其仅包含当前可用的对象路径。

确保您记录已删除的 compat.v1.layers 的范围，以保留为剩余的 compat.v1.layers 自动生成的权重名称。

weights = set(v.name for v in layer.weights)
assert 'model/conv2d_2/kernel:0' in weights
assert 'model/conv2d_2/bias:0' in weights

1. 重复上述步骤，直到您将模型中的所有 compat.v1.layers 和 compat.v1.get_variable 替换为完全的原生等效项。

class FullyNativeKerasLayersModel(tf.keras.layers.Layer):

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.conv_1 = tf.keras.layers.Conv2D(
          3, 3,
          kernel_regularizer="l2")
    self.conv_2 = tf.keras.layers.Conv2D(
          4, 4,
          kernel_regularizer="l2")
    self.conv_3 = tf.keras.layers.Conv2D(
          5, 5,
          kernel_regularizer="l2")


  def call(self, inputs, training=None):
    with tf.compat.v1.variable_scope('model'):
      out = self.conv_1(inputs)
      out = self.conv_2(out)
      out = self.conv_3(out)
      return out

layer = FullyNativeKerasLayersModel()
layer(tf.ones(shape=(10, 10, 10, 10)))
[v.name for v in layer.weights]

请记住测试以确保新更新的检查点仍然按预期运行。在该过程的每个增量步骤中应用 验证数值正确性指南 中描述的技术，以确保您迁移的代码正确运行。

处理模型屏蔽未涵盖的 TF1.x 到 TF2 行为更改

本指南中描述的模型屏蔽可以确保使用 get_variable、tf.compat.v1.layers 和 variable_scope 语义创建的变量、层和正则化损失在使用急切执行和 tf.function 时继续按预期工作，而无需依赖集合。

这并不涵盖您的模型前向传递可能依赖的所有 TF1.x 特定语义。在某些情况下，屏蔽可能不足以使您的模型前向传递仅靠它们本身在 TF2 中运行。阅读 TF1.x 与 TF2 行为指南 以了解有关 TF1.x 和 TF2 之间行为差异的更多信息。