8. Tensorflow中,环境变量TF_CONFIG如何利用已知变量进行构建?

以Tensorflow Estimator分布式中,chief模式下的环境变量TF_CONFIG的构建为例(详细代码可见$XLEARNING_HOME/examples/tfEstimator),如下:

  1. import os
  2. import json
  3. cluster = json.loads(os.environ["TF_CLUSTER_DEF"])
  4. task_index = int(os.environ["TF_INDEX"])
  5. task_type = os.environ["TF_ROLE"]
  6. # chief: worker 0 as chief, other worker index --
  7. tf_config = dict()
  8. worker_num = len(cluster["wroker"])
  9. if task_type == "ps":
  10. tf_config["task"] = {"index":task_index, "type":task_type}
  11. elif task_type == "worker":
  12. if taks_index == 0:
  13. tf_config["task"] = {"index":0, "type":"chief"}
  14. else:
  15. tf_config["task"] = {"index":task_index-1, "type":task_type}
  16. elif task_type == "evaluator":
  17. tf_config["task"] = {"index":task_index, "type":task_type}
  18. if worker_num == 1:
  19. cluster["chief"] = cluster["worker"]
  20. del cluster["worker"]
  21. else:
  22. cluster["chief"] = [cluster["worker"][0]]
  23. del cluster["worker"][0]
  24. tf_config["cluster"] = cluster
  25. os.environ["TF_CONFIG"] = json.dumps(tf_config)

由此,可利用Tensorflow分布式模式下,XLearning提供的环境变量 TF_CLUSTER_DEF 、 TF_ROLE 、 TF_INDEX 对应的来构建所需的环境变量TF_CONFIG。