Tensorflow Estimator源码分析

Estimator是Tensorflow的高阶API。除了Tensorflow官方定义的内置Estimator之外，用户也可以实现自定义的Estimator。

Estimator定义

Estimator的构造函数如下：

  def __init__(self,
    model_fn,  # 定义模型，根据不同的模式分别定义训练、评估和预测的图。
    model_dir=None,  # 模型导出目录
    config=None,     # 配置参数
    params=None,     # 自定义Estimator的额外参数
    warm_start_from=None):  # 模型热启动

其中最核心的参数为model_fn，其接口如下

  def _model_fn(features,  # 特征，可以是Tensor或dict of Tensor
                labels,    # 标签
                mode,      # 模式
                params,    # 自定义参数，即上面Estimator构造函数中的params
                config):   # 配置参数

model_fn会被Estimator多次调用，通过调用Tensorflow的layer来实现模型。通过模式字段(ModeKeys.TRAIN, ModeKeys.EVAL, ModeKeys.PREDICT)来判断是训练、评估还是预测阶段，分别构造不同的图。model_fn的返回结构为EstimatorSpec，使用其中的训练、loss和预测的OP，Estimator就可以驱动完成训练、评估和预测。

EstimatorSpec的定义如下

  def __new__(cls,
              mode,  # 模式
              predictions=None,  # 预测的Tensor或dict，mode为PREDICT时必填。
              loss=None,  # loss Tensor，mode为TRAIN或EVAL时必填。
              train_op=None,  # 训练OP，mode为TRAIN时必填。
              eval_metric_ops=None,  # 评估OP的dict
              export_outputs=None,
              training_chief_hooks=None,
              training_hooks=None,
              scaffold=None,
              evaluation_hooks=None,
              prediction_hooks=None):

训练

Estimator的训练接口如下

  def train(self,
            input_fn,    # 返回训练特征和标签的tuple
            hooks=None,  # 通过hook指定训练过程中的自定义行为
            steps=None,  # 训练步数
            max_steps=None,  ## 训练总步数
            saving_listeners=None):
    with context.graph_mode():
      hooks.extend(self._convert_train_steps_to_hooks(steps, max_steps))
      loss = self._train_model(input_fn, hooks, saving_listeners)
      logging.info('Loss for final step: %s.', loss)

_train_model根据不同的配置，分别走到分布式训练和本地训练的函数。

  def _train_model(self, input_fn, hooks, saving_listeners):
    if self._train_distribution:
      return self._train_model_distributed(input_fn, hooks, saving_listeners)
    else:
      return self._train_model_default(input_fn, hooks, saving_listeners)

我们先看本地训练的实现。

  def _train_model_default(self, input_fn, hooks, saving_listeners):
    with ops.Graph().as_default() as g, g.device(self._device_fn):
      random_seed.set_random_seed(self._config.tf_random_seed)
      global_step_tensor = self._create_and_assert_global_step(g)
      features, labels, input_hooks = (
          self._get_features_and_labels_from_input_fn(
              input_fn, ModeKeys.TRAIN))
      worker_hooks.extend(input_hooks)
      estimator_spec = self._call_model_fn(
          features, labels, ModeKeys.TRAIN, self.config)
      global_step_tensor = training_util.get_global_step(g)
      return self._train_with_estimator_spec(estimator_spec, worker_hooks,
                                             hooks, global_step_tensor,
                                             saving_listeners)

其流程为先创建global_step，然后调用input_fn来得到训练特征和标签，调用model_fn来得到训练图，最后进入training loop。

_get_features_and_labels_from_input_fn最终会调用input_fn，得到训练特征和标签。

    with ops.device('/cpu:0'):
      return input_fn(**kwargs)

_call_model_fn会调用model_fn，注意传递的参数为ModeKeys.TRAIN，用于表征训练阶段。

def _call_model_fn(self, features, labels, mode, config):
    model_fn_results = self._model_fn(features=features, **kwargs)

下面看_train_with_estimator_spec的实现。

  def _train_with_estimator_spec(self, estimator_spec, worker_hooks, hooks,
                                 global_step_tensor, saving_listeners):
    # 满足条件则热启动
    if self._warm_start_settings:
      warm_starting_util.warm_start(*self._warm_start_settings)
    # 创建Hook
    worker_hooks.extend(hooks)
    worker_hooks.append(training.NanTensorHook(estimator_spec.loss)
    worker_hooks.append(training.LoggingTensorHook(...))
    saver_hooks = [
        h for h in all_hooks if isinstance(h, training.CheckpointSaverHook)]
    worker_hooks.extend(estimator_spec.training_hooks)
    worker_hooks.append(training.SummarySaverHook(...))
    worker_hooks.append(training.StepCounterHook(...))

    with training.MonitoredTrainingSession(
        master=self._config.master,
        is_chief=self._config.is_chief,
        checkpoint_dir=self._model_dir,
        scaffold=estimator_spec.scaffold,
        hooks=worker_hooks,
        chief_only_hooks=(
            tuple(chief_hooks) + tuple(estimator_spec.training_chief_hooks)),
        save_checkpoint_secs=0,  # Saving is handled by a hook.
        save_summaries_steps=save_summary_steps,
        config=self._session_config,
        log_step_count_steps=log_step_count_steps) as mon_sess:
      loss = None
      any_step_done = False
      while not mon_sess.should_stop():
        _, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])
        any_step_done = True
    if not any_step_done:
      logging.warning('Training with estimator made no steps. '
                      'Perhaps input is empty or misspecified.')
    return loss

前面主要在创建Hook，后面使用MonitoredTrainingSession进行Training loop。

评估

评估的接口为

  def evaluate(self, input_fn, steps=None, hooks=None, checkpoint_path=None,
               name=None):

其中input_fn接口与训练函数中的input_fn有相同的接口，调用后返回评估用的特征和标签。评估最终会调用到下面的函数

  def _actual_eval(self,
                   input_fn,
                   strategy=None,
                   steps=None,
                   hooks=None,
                   checkpoint_path=None,
                   name=None):
      ...
      def _evaluate():
        (scaffold, update_op, eval_dict, all_hooks) = (
            self._evaluate_build_graph(input_fn, hooks, checkpoint_path))
        return self._evaluate_run(
            checkpoint_path=checkpoint_path,
            scaffold=scaffold,
            update_op=update_op,
            eval_dict=eval_dict,
            all_hooks=all_hooks,
            output_dir=self.eval_dir(name))
      return _evaluate()

_evaluate_build_graph的实现如下：

  def _evaluate_build_graph(self, input_fn, hooks=None, checkpoint_path=None):
    """Builds the graph and related hooks to run evaluation."""
    (scaffold, evaluation_hooks, input_hooks, update_op, eval_dict) = (
          self._call_model_fn_eval(input_fn, self.config))

    all_hooks = list(input_hooks)
    all_hooks.extend(hooks)
    all_hooks.extend(list(evaluation_hooks or []))
    if scaffold and scaffold.local_init_op:
      # 创建评估step
      evaluation._get_or_create_eval_step()  # pylint: disable=protected-access

      scaffold = monitored_session.Scaffold(
          local_init_op=control_flow_ops.group(
              scaffold.local_init_op,
              monitored_session.Scaffold.default_local_init_op()),
          copy_from_scaffold=scaffold
      )
    return scaffold, update_op, eval_dict, all_hooks

_evaluate_build_graph会调用_call_model_fn_eval，进行评估构图，然后返回scaffold。

  def _call_model_fn_eval(self, input_fn, config):
    """Call model_fn for evaluation and handle return values."""
    features, labels, input_hooks = self._get_features_and_labels_from_input_fn(
        input_fn, ModeKeys.EVAL)

    estimator_spec = self._call_model_fn(
        features, labels, ModeKeys.EVAL, config)
    eval_metric_ops = _verify_and_create_loss_metric(
        estimator_spec.eval_metric_ops, estimator_spec.loss)
    update_op, eval_dict = _extract_metric_update_ops(eval_metric_ops)
    return (estimator_spec.scaffold, estimator_spec.evaluation_hooks,
            input_hooks, update_op, eval_dict)

_call_model_fn_eval流程为从input_fn获取评估用的特征和标签，然后调用model_fn进行评估构图。
_actual_eval调用完_evaluate_build_graph之后，接着调用_evaluate_run。

  def _evaluate_run(self, checkpoint_path, scaffold, update_op, eval_dict,
                    all_hooks, output_dir):
    """Run evaluation."""
    eval_results = evaluation._evaluate_once(  # pylint: disable=protected-access
        checkpoint_path=checkpoint_path,
        master=self._config.evaluation_master,
        scaffold=scaffold,
        eval_ops=update_op,
        final_ops=eval_dict,
        hooks=all_hooks,
        config=self._session_config)
    ...

def _evaluate_once(checkpoint_path,
                   master='',
                   scaffold=None,
                   eval_ops=None,
                   feed_dict=None,
                   final_ops=None,
                   final_ops_feed_dict=None,
                   hooks=None,
                   config=None):
    # 准备eval_ops
    if isinstance(eval_ops, dict):
      eval_ops['update_eval_step'] = update_eval_step
    elif isinstance(eval_ops, (tuple, list)):
      eval_ops = list(eval_ops) + [update_eval_step]
    else:
      eval_ops = [eval_ops, update_eval_step]

    eval_step_value = _get_latest_eval_step_value(eval_ops)

  # Prepare the session creator.
  session_creator = monitored_session.ChiefSessionCreator(
      scaffold=scaffold,
      checkpoint_filename_with_path=checkpoint_path,
      master=master,
      config=config)

  with monitored_session.MonitoredSession(
      session_creator=session_creator, hooks=hooks) as session:
    if eval_ops is not None:
      while not session.should_stop():
        session.run(eval_ops, feed_dict)

_evaluate_once执行最终的评估逻辑，先准备好评估用的ops，然后通过MonitoredSession执行评估的loop。

预测

预测的接口和实现如下，相对最为简单。

  def predict(self,
              input_fn,
              predict_keys=None,
              hooks=None,
              checkpoint_path=None,
              yield_single_examples=True):
      with ops.Graph().as_default() as g:
        # 从`input_fn`获取预测用的特征。
        features, input_hooks = self._get_features_from_input_fn(
            input_fn, ModeKeys.PREDICT)
        estimator_spec = self._call_model_fn(
            features, None, ModeKeys.PREDICT, self.config)

        predictions = self._extract_keys(
            estimator_spec.predictions, predict_keys)
        with training.MonitoredSession(
            session_creator=training.ChiefSessionCreator(
                checkpoint_filename_with_path=checkpoint_path,
                master=self._config.master,
                scaffold=estimator_spec.scaffold,
                config=self._session_config),
            hooks=all_hooks) as mon_sess:
          while not mon_sess.should_stop():
            preds_evaluated = mon_sess.run(predictions)

导出模型

Estimator最后一个重要接口为导出模型接口，

  def export_saved_model(
      self, export_dir_base, serving_input_receiver_fn,
      assets_extra=None,
      as_text=False,
      checkpoint_path=None,
      experimental_mode=ModeKeys.PREDICT):
    input_receiver_fn_map = {experimental_mode: serving_input_receiver_fn}
    return self._export_all_saved_models(
        export_dir_base,
        input_receiver_fn_map,
        assets_extra=assets_extra,
        as_text=as_text,
        checkpoint_path=checkpoint_path,
        strip_default_attrs=True)

  def _export_all_saved_models(
      self, export_dir_base, input_receiver_fn_map,
      assets_extra=None, as_text=False, checkpoint_path=None,
      strip_default_attrs=True):
    with context.graph_mode():
      builder = saved_model_builder.SavedModelBuilder(temp_export_dir)
      if input_receiver_fn_map.get(ModeKeys.PREDICT):
        self._add_meta_graph_for_mode(
            builder, input_receiver_fn_map, checkpoint_path,
            save_variables, mode=ModeKeys.PREDICT,
            strip_default_attrs=strip_default_attrs)
      builder.save(as_text)

内置Estimator

我们看一下LinearClassifierV2的实现

class LinearClassifierV2(estimator.EstimatorV2):
  def __init__(self,
               feature_columns,
               model_dir=None,
               n_classes=2,
               weight_column=None,
               label_vocabulary=None,
               optimizer='Ftrl',
               config=None,
               warm_start_from=None,
               loss_reduction=losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
               sparse_combiner='sum'):
    head = head_utils.binary_or_multi_class_head(
        n_classes, weight_column=weight_column,
        label_vocabulary=label_vocabulary,
        loss_reduction=loss_reduction)

    def _model_fn(features, labels, mode, config):
      """Call the defined shared _linear_model_fn."""
      return _linear_model_fn_v2(
          features=features,
          labels=labels,
          mode=mode,
          head=head,
          feature_columns=tuple(feature_columns or []),
          optimizer=optimizer,
          config=config,
          sparse_combiner=sparse_combiner)

    super(LinearClassifierV2, self).__init__(
        model_fn=_model_fn,
        model_dir=model_dir,
        config=config,
        warm_start_from=warm_start_from)

可以看到内置Estimator的实现和自定义Estimator的实现没什么区别，也是通过实现model_fn并创建Estimator实例得到的。