- 8. Tensorflow中,环境变量TF_CONFIG如何利用已知变量进行构建?
8. Tensorflow中,环境变量TF_CONFIG如何利用已知变量进行构建?
以Tensorflow Estimator分布式中,chief模式下的环境变量TF_CONFIG的构建为例(详细代码可见$XLEARNING_HOME/examples/tfEstimator),如下:
import os
import json
cluster = json.loads(os.environ["TF_CLUSTER_DEF"])
task_index = int(os.environ["TF_INDEX"])
task_type = os.environ["TF_ROLE"]
# chief: worker 0 as chief, other worker index --
tf_config = dict()
worker_num = len(cluster["wroker"])
if task_type == "ps":
tf_config["task"] = {"index":task_index, "type":task_type}
elif task_type == "worker":
if taks_index == 0:
tf_config["task"] = {"index":0, "type":"chief"}
else:
tf_config["task"] = {"index":task_index-1, "type":task_type}
elif task_type == "evaluator":
tf_config["task"] = {"index":task_index, "type":task_type}
if worker_num == 1:
cluster["chief"] = cluster["worker"]
del cluster["worker"]
else:
cluster["chief"] = [cluster["worker"][0]]
del cluster["worker"][0]
tf_config["cluster"] = cluster
os.environ["TF_CONFIG"] = json.dumps(tf_config)
由此,可利用Tensorflow分布式模式下,XLearning提供的环境变量 TF_CLUSTER_DEF 、 TF_ROLE 、 TF_INDEX 对应的来构建所需的环境变量TF_CONFIG。