task.py 文件源码

def main(argv=None):
  """Run a Tensorflow model on the Criteo dataset."""
  env = json.loads(os.environ.get('TF_CONFIG', '{}'))
  # First find out if there's a task value on the environment variable.
  # If there is none or it is empty define a default one.
  task_data = env.get('task') or {'type': 'master', 'index': 0}
  argv = sys.argv if argv is None else argv
  args = create_parser().parse_args(args=argv[1:])

  trial = task_data.get('trial')
  if trial is not None:
    output_dir = os.path.join(args.output_path, trial)
  else:
    output_dir = args.output_path

  # Do only evaluation if instructed so, or call Experiment's run.
  if args.eval_only_summary_filename:
    experiment = get_experiment_fn(args)(output_dir)
    # Note that evaluation here will appear as 'one_pass' in tensorboard.
    results = experiment.evaluate(delay_secs=0)
    # Converts numpy types to native types for json dumps.
    json_out = json.dumps(
        {key: value.tolist() for key, value in results.iteritems()})
    with tf.Session():
      tf.write_file(args.eval_only_summary_filename, json_out).run()
  else:
    learn_runner.run(experiment_fn=get_experiment_fn(args),
                     output_dir=output_dir)