远离 tf.contrib.learn：具有专用评估流程的分布式培训答案

【问题标题】：Moving away from tf.contrib.learn: distributed training with dedicated evaluator process远离 tf.contrib.learn：具有专用评估流程的分布式培训
【发布时间】：2018-04-24 17:10:20
【问题描述】：

在 TF 1.8 即将发布的版本中，tf.contrib.learn.* 将被弃用。 tf.contrib.learn.Experiment 类建议改用 tf.estimator.train_and_evaluate，因此我正在尝试将我的代码移植到该框架。

我想做的是在两台机器的 GPU 上设置分布式训练，再加上第三个仅 CPU 的进程，该进程对小型验证集进行连续评估。

按照the documentation of train_and_evaluate 和Distributed Tensorflow 指南中的示例，我设法设置了所需架构的训练部分，但我找不到设置估算器的方法。

到目前为止，我看到的内容如下：

def input_fn(mode, num_classes, batch_size):  
  # [...] build input pipeline
  return {'input': images}, labels

def model_fn(features, labels, num_classes, mode):
  # [...] build model
  return tf.estimator.EstimatorSpec(
    mode=mode,
    predictions=predictions,
    loss=total_loss,
    train_op=train_op,
    eval_metric_ops=metrics,
    export_outputs=export_outputs)

def distributed_main_v2(unused_argv):
  """Expects `unused_argv` to be a list ['<task_type>', '<task_id>']"""  
  import json
  # Set up environment variables according to the parameters passed to the process
  TF_CONFIG = {
    'cluster': {
        "ps": [
            "host1:2222",
        ],
        "chief": [
            "host1:2223",
            ],
        "worker": [
            "host2:2224"
            ]
    },
    'environment': 'cluster',    
    'task': {
        'type': unused_argv[1].strip(),
        'index': unused_argv[2].strip() if len(unused_argv) > 2 else 0
        }
  }
  os.environ['TF_CONFIG'] = json.dumps(TF_CONFIG)
  if unused_argv[1].strip() not in ['worker', 'chief']:
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # leave the GPU to the worker process

  # create the estimator
  # define warm start configuration
  regex = '^(?!.*final_layer*|.*aux_logits*)'
  ws_settings = tf.estimator.WarmStartSettings('checkpoint_path', regex)

  gpu_opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.95) # fix for cuDNN fatal memory error with tf.contrib.learn.Experiment (TODO: still necessary?)
  sess_conf = tf.ConfigProto(gpu_options=gpu_opts)
  run_conf = tf.estimator.RunConfig(session_config=sess_conf)

  # Create the Estimator
  estimator = tf.estimator.Estimator(
    model_fn=lambda features, labels, mode: model_fn(features, labels, NUM_CLASSES, mode),
    model_dir=model_dir,
    config=run_conf,
    warm_start_from=ws_settings)

  # Set up input functions for training and evaluation
  train_input_fn = lambda : input_fn(tf.estimator.ModeKeys.TRAIN, NUM_CLASSES, batch_size)
  eval_input_fn = lambda : input_fn(tf.estimator.ModeKeys.EVAL, NUM_CLASSES, batch_size)

  train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=steps)
  eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)

  # start distributed training
  tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

if __name__ == '__main__':
  # set up globals and parse known arguments
  distributed_main_v2(unused_argv)

此代码有效，尽管我对它的理解仍然非常有限。我得到了 PS 和工作人员所做的事情，但从 chief 的规范中，我知道这应该是“主”工作人员，它也记录摘要并保存检查点。我现在缺少的是定期评估……我很茫然。从train_and_evaluate 代码库我看到有一些“评估器”支持，但我不明白如何正确设置它。

【问题讨论】：

标签： python tensorflow

【解决方案1】：

注意：在写这个问题时，我最终意识到了我的错误（即，我是盲目的，没有看到到目前为止我认为我至少看了 20 次的代码和文档），但我相信问题和匹配的答案可能对其他人有用，所以我决定完成问题并自行回答。

如果我阅读整个docs 的内容，我会注意到以下内容：

评估任务的 TF_CONFIG 示例。评估员是一项特殊任务这不是训练集群的一部分。可能只有一个。它用于模型评估。

# This should be a JSON string, which is set as environment variable. Usually
# the cluster manager handles that.
TF_CONFIG='{
    "cluster": {
        "chief": ["host0:2222"],
        "worker": ["host1:2222", "host2:2222", "host3:2222"],
        "ps": ["host4:2222", "host5:2222"]
    },
    "task": {"type": "evaluator", "index": 0}
}'

事实证明，是的，确实支持评估任务并且使用它比我预期的要容易得多。

只需将TF_CONFIG 的"task"part 设置为{"type": "evaluator", "index": 0}，如上所示，您就可以运行评估了。让我感到困惑的是“评估器是一项不属于训练集群的特殊任务”。我相信这是因为首席工作人员在开始分布式会话时等待所有工作人员向他注册，因此将评估者排除在集群之外可以保持训练和评估彼此独立，并使训练与评估无关。

【讨论】：