"git@developer.sourcefind.cn:wangsen/mineru.git" did not exist on "d2228fb93fb6e7b138e65fcac21485a15685ec7f"
Unverified Commit 9f7a5fa3 authored by Taylor Robie's avatar Taylor Robie Committed by GitHub
Browse files

Replicate internal cluster changes in official mnist tpu example (#4020)

parent 26aded58
......@@ -46,10 +46,6 @@ tf.flags.DEFINE_string(
"metadata.")
# Model specific parameters
tf.flags.DEFINE_string(
"master", default=None,
help="GRPC URL of the master (e.g. grpc://ip.address.of.tpu:8470). You "
"must specify either this flag or --tpu.")
tf.flags.DEFINE_string("data_dir", "",
"Path to directory containing the MNIST dataset")
tf.flags.DEFINE_string("model_dir", None, "Estimator model_dir")
......@@ -136,24 +132,14 @@ def main(argv):
del argv # Unused.
tf.logging.set_verbosity(tf.logging.INFO)
if FLAGS.master is None and FLAGS.tpu is None:
raise RuntimeError('You must specify either --master or --tpu.')
if FLAGS.master is not None:
if FLAGS.tpu is not None:
tf.logging.warn('Both --master and --tpu are set. Ignoring '
'--tpu and using --master.')
tpu_grpc_url = FLAGS.master
else:
tpu_cluster_resolver = (
tf.contrib.cluster_resolver.TPUClusterResolver(
FLAGS.tpu,
zone=FLAGS.tpu_zone,
project=FLAGS.gcp_project))
tpu_grpc_url = tpu_cluster_resolver.get_master()
tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
FLAGS.tpu,
zone=FLAGS.tpu_zone,
project=FLAGS.gcp_project
)
run_config = tf.contrib.tpu.RunConfig(
master=tpu_grpc_url,
evaluation_master=tpu_grpc_url,
cluster=tpu_cluster_resolver,
model_dir=FLAGS.model_dir,
session_config=tf.ConfigProto(
allow_soft_placement=True, log_device_placement=True),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment