從 Estimator 遷移至 Keras API

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

本指南示範如何從 TensorFlow 1 的 tf.estimator.Estimator API 遷移至 TensorFlow 2 的 tf.keras API。首先，您將設定並執行基本模型，以便使用 tf.estimator.Estimator 進行訓練和評估。接著，您將在 TensorFlow 2 中使用 tf.keras API 執行同等步驟。您也將學習如何透過子類別化 tf.keras.Model 和使用 tf.GradientTape 來自訂訓練步驟。

  • 在 TensorFlow 1 中，高階 tf.estimator.Estimator API 可讓您訓練和評估模型，以及執行推論並儲存模型 (以供服務使用)。
  • 在 TensorFlow 2 中，使用 Keras API 執行上述工作，例如 模型建構、梯度應用、訓練、評估和預測。

(如需將模型/檢查點儲存工作流程遷移至 TensorFlow 2，請參閱 SavedModel 和 Checkpoint 遷移指南。)

設定

從匯入項目和簡單資料集開始

import tensorflow as tf
import tensorflow.compat.v1 as tf1

features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [[0.8], [0.9], [1.]]

TensorFlow 1：使用 tf.estimator.Estimator 進行訓練和評估

本範例說明如何在 TensorFlow 1 中使用 tf.estimator.Estimator 執行訓練和評估。

首先定義幾個函式：訓練資料的輸入函式、評估資料的評估輸入函式，以及模型函式，告知 Estimator 如何使用特徵和標籤定義訓練運算

def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

def _eval_input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

def _model_fn(features, labels, mode):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

例項化您的 Estimator，並訓練模型

estimator = tf1.estimator.Estimator(model_fn=_model_fn)
estimator.train(_input_fn)

使用評估集評估程式

estimator.evaluate(_eval_input_fn)

TensorFlow 2：使用內建 Keras 方法進行訓練和評估

本範例示範如何在 TensorFlow 2 中使用 Keras Model.fit 和 Model.evaluate 執行訓練和評估。(您可以在「使用內建方法進行訓練和評估」指南中瞭解更多資訊。)

  • 首先使用 tf.data.Dataset API 準備資料集管線。
  • 定義具有一個線性 (tf.keras.layers.Dense) 層的簡單 Keras Sequential 模型。
  • 例項化 Adagrad 最佳化工具 (tf.keras.optimizers.Adagrad)。
  • 透過將 optimizer 變數和均方誤差 ("mse") 損失傳遞至 Model.compile，設定模型以進行訓練。
dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss="mse")

如此一來，您就可以透過呼叫 Model.fit 來訓練模型

model.fit(dataset)

最後，使用 Model.evaluate 評估模型

model.evaluate(eval_dataset, return_dict=True)

TensorFlow 2：使用自訂訓練步驟和內建 Keras 方法進行訓練和評估

在 TensorFlow 2 中，您也可以使用 tf.GradientTape 撰寫自己的自訂訓練步驟函式，以執行正向和反向傳播，同時仍可利用內建訓練支援，例如 tf.keras.callbacks.Callback 和 tf.distribute.Strategy。(如需更多資訊，請參閱「自訂 Model.fit 中的事件」和「從頭開始撰寫自訂訓練迴圈」。)

在本範例中，首先透過子類別化 tf.keras.Sequential 來建立自訂 tf.keras.Model，以覆寫 Model.train_step。(深入瞭解子類別化 tf.keras.Model)。在該類別中，定義自訂 train_step 函式，該函式會在一個訓練步驟中針對每個資料批次執行正向傳播和反向傳播。

class CustomModel(tf.keras.Sequential):
  """A custom sequential model that overrides `Model.train_step`."""

  def train_step(self, data):
    batch_data, labels = data

    with tf.GradientTape() as tape:
      predictions = self(batch_data, training=True)
      # Compute the loss value (the loss function is configured
      # in `Model.compile`).
      loss = self.compiled_loss(labels, predictions)

    # Compute the gradients of the parameters with respect to the loss.
    gradients = tape.gradient(loss, self.trainable_variables)
    # Perform gradient descent by updating the weights/parameters.
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
    # Update the metrics (includes the metric that tracks the loss).
    self.compiled_metrics.update_state(labels, predictions)
    # Return a dict mapping metric names to the current values.
    return {m.name: m.result() for m in self.metrics}

接下來，如同先前

dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

model = CustomModel([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss="mse")

呼叫 Model.fit 以訓練模型

model.fit(dataset)

最後，使用 Model.evaluate 評估程式

model.evaluate(eval_dataset, return_dict=True)

後續步驟

您可能會覺得實用的其他 Keras 資源

以下指南可協助從 tf.estimator API 遷移分散式策略工作流程