aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2018-09-25 13:06:23 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-09-25 13:17:25 -0700
commit22c0efbc82d4c721d7a96b7acc1205e161563800 (patch)
tree2baa0bcf3cf9e42742e2c128f29b63f30e549ef5
parenta50dff24b6f38fef7ead20e1015509cac905ed29 (diff)
Optimize Resnet-50 infeed pipeline by adding sharding and cache.
PiperOrigin-RevId: 214489904
-rw-r--r--tensorflow/contrib/tpu/python/tpu/tpu_context.py5
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py
index b1a8a16d1e..a5f86ecd5f 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_context.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py
@@ -118,6 +118,11 @@ class TPUContext(object):
return self._internal_ctx.num_hosts
@property
+ def current_host(self):
+ """The current host index for the TPU system."""
+ return self._invocation_index
+
+ @property
def num_of_replicas_per_host(self):
"""The number of replicas for each host."""
if self._internal_ctx.model_parallelism_enabled: