从 Estimator 迁移到 Keras API

本指南演示了如何从 TensorFlow 1 的 tf.estimator.Estimator API 迁移到 TensorFlow 2 的 tf.keras API。首先，您将设置并运行一个基本的模型，以使用 tf.estimator.Estimator 进行训练和评估。然后，您将使用 tf.keras API 在 TensorFlow 2 中执行等效步骤。您还将学习如何通过子类化 tf.keras.Model 并使用 tf.GradientTape 来定制训练步骤。

在 TensorFlow 1 中，高级 tf.estimator.Estimator API 允许您训练和评估模型，以及执行推理和保存模型（用于服务）。
在 TensorFlow 2 中，使用 Keras API 执行上述任务，例如模型构建、梯度应用、训练、评估和预测。

（要将模型/检查点保存工作流程迁移到 TensorFlow 2，请查看 SavedModel 和检查点迁移指南。）

设置

从导入和简单数据集开始

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}

接下来，和之前一样

使用 tf.data.Dataset 准备数据集管道。
使用一个 tf.keras.layers.Dense 层定义一个简单的模型。
实例化 Adagrad (tf.keras.optimizers.Adagrad)
使用 Model.compile 配置模型进行训练，同时使用均方误差 ("mse") 作为损失函数。

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 迁移