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

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視 下載筆記本

本指南概述並舉例說明您可以採用的模型程式碼墊片,以便在 TF2 工作流程 (例如立即執行、tf.function 和分散式策略) 中使用現有的 TF1.x 模型,且對模型程式碼的變更最少。

使用範圍

本指南中描述的墊片專為依賴下列項目的 TF1.x 模型而設計:

  1. tf.compat.v1.get_variabletf.compat.v1.variable_scope,以控制變數建立和重複使用,以及
  2. 以圖形集合為基礎的 API,例如 tf.compat.v1.global_variables()tf.compat.v1.trainable_variablestf.compat.v1.losses.get_regularization_losses()tf.compat.v1.get_collection(),以追蹤權重和正規化損失

這包括大多數建構於 tf.compat.v1.layertf.contrib.layers API 和 TensorFlow-Slim 之上的模型。

下列 TF1.x 模型需要墊片

  1. 獨立的 Keras 模型,這些模型已分別透過 model.trainable_weightsmodel.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.Layertf.Module 的方法中使用它,以追蹤 TF1.x 樣式的權重並擷取正規化損失。

使用 tf.compat.v1.keras.utils.track_tf1_style_variables 裝飾 tf.keras.layers.Layertf.Module 的呼叫方法,可讓透過 tf.compat.v1.get_variable (以及擴充功能 tf.compat.v1.layers) 建立和重複使用變數,在裝飾方法內正常運作,而不是始終在每次呼叫時建立新變數。它也會讓圖層或模組隱含追蹤在裝飾方法內透過 get_variable 建立或存取的任何權重。

除了在標準 layer.variable/module.variable 等屬性下追蹤權重本身之外,如果該方法屬於 tf.keras.layers.Layer,則透過 get_variabletf.compat.v1.layers 正規化工具引數指定的任何正規化損失,都會由圖層在標準 layer.losses 屬性下追蹤。

這種追蹤機制讓大型 TF1.x 樣式的模型正向傳遞程式碼可以在 TF2 中的 Keras 圖層或 tf.Module 內使用,即使啟用 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.Layertf.Module 的裝飾方法會在立即執行內執行,並且也與 tf.function 相容。這表示您可以使用 pdb 和其他互動式工具逐步執行正在執行的正向傳遞。

分散式策略

@track_tf1_style_variables 裝飾的圖層或模組方法內呼叫 get_variable 會在底層使用標準 tf.Variable 變數建立。這表示您可以將它們與 tf.distribute 提供的各種分散式策略搭配使用,例如 MirroredStrategyTPUStrategy

在裝飾的呼叫中巢狀 tf.Variabletf.Moduletf.keras.layerstf.keras.models

tf.compat.v1.keras.utils.track_tf1_style_variables 中裝飾圖層呼叫只會新增透過 tf.compat.v1.get_variable 建立 (和重複使用) 的變數的自動隱含追蹤。它不會擷取由 tf.Variable 呼叫直接建立的權重,例如一般 Keras 圖層和大多數 tf.Module 使用的權重。本節說明如何處理這些巢狀案例。

(先前用法) tf.keras.layerstf.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.Variablestf.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_variablevariable_scopes 的組合中習慣使用的機制不同。雖然墊片會讓您的變數名稱符合由 get_variable 建立的變數,即使從 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 用法逐步切換為建議用於新撰寫的 TF2 程式碼的純原生物件導向 API。

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 更複雜。實際上,將 Slim 程式碼轉換為 v1.layers,然後再轉換為 Keras 可能更有意義。以下是轉換 Slim 程式碼的一些一般指南。

  • 確保所有引數都是明確的。如果可能,請移除 arg_scopes。如果您仍然需要使用它們,請將 normalizer_fnactivation_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 物件的屬性,其方法以 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.layers 保留變數名稱,方法是手動開啟和關閉對應於已移除 compat.v1.layer 生成範圍名稱的 variable_scope。否則,來自現有檢查點的路徑可能會衝突,並且檢查點載入的行為將不正確。
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.layerscompat.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_variabletf.compat.v1.layersvariable_scope 語意建立的變數、層和正規化損失,在使用迫切執行和 tf.function 時,繼續像以前一樣工作,而無需依賴集合。

這不涵蓋您的模型正向傳遞可能依賴的所有 TF1.x 特定語意。在某些情況下,墊片可能不足以讓您的模型正向傳遞在 TF2 中自行執行。請閱讀 TF1.x 與 TF2 行為指南,以瞭解有關 TF1.x 和 TF2 之間行為差異的更多資訊。