从 Estimator 迁移到 Keras API

##### Copyright 2021 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

从 Estimator 迁移到 Keras API#

在 TensorFlow.org 上查看 在 Google Colab 运行 在 Github 上查看源代码 下载笔记本

本指南演示了如何从 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 中使用 Model.fit 和 Model.evaluate 执行训练和评估。（可以在使用内置方法进行训练和评估指南中了解详情。）

  • 首先使用 tf.data.Dataset API 准备数据集流水线。

  • 使用一个线性 (tf.keras.layers.Dense) 层定义一个简单的 Keras 序贯模型。

  • 实例化一个 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 的功能和从头开始编写自定义训练循环中了解详情。）

在此示例中，首先通过子类化重写 Model.train_step 的 tf.keras.Sequential 来创建自定义 tf.keras.Model。（详细了解如何子类化 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 迁移分布策略工作流：