diff options
Diffstat (limited to 'tensorflow/python/data/experimental/ops/shuffle_ops.py')
| -rw-r--r-- | tensorflow/python/data/experimental/ops/shuffle_ops.py | 21 |
1 files changed, 18 insertions, 3 deletions
diff --git a/tensorflow/python/data/experimental/ops/shuffle_ops.py b/tensorflow/python/data/experimental/ops/shuffle_ops.py index a4307212da..a82e4b7d09 100644 --- a/tensorflow/python/data/experimental/ops/shuffle_ops.py +++ b/tensorflow/python/data/experimental/ops/shuffle_ops.py @@ -39,17 +39,32 @@ class _ShuffleAndRepeatDataset(dataset_ops.UnaryDataset): else: self._count = ops.convert_to_tensor( count, dtype=dtypes.int64, name="count") - self._seed, self._seed2 = random_seed.get_seed(seed) + + # NOTE(mrry): We generate the seed-pair once per graph in which the dataset + # is iterated over, and cache it in `self._graph_seed_map`. This supports + # two features: iterating over the same `ShuffleDataset` twice in the same + # pipeline and observing the same order (by tying the seeds together with + # a randomly-generated seed), and using `Dataset.make_one_shot_iterator()`, + # which requires the stateful RNG op to be created inside the same graph as + # the dataset. + self._original_seed = seed + self._graph_seed_map = {} def _as_variant_tensor(self): + try: + seed, seed2 = self._graph_seed_map[ops.get_default_graph()] + except KeyError: + seed, seed2 = random_seed.get_seed(self._original_seed) + self._graph_seed_map[ops.get_default_graph()] = (seed, seed2) + # pylint: disable=protected-access input_resource = self._input_dataset._as_variant_tensor() return gen_dataset_ops.shuffle_and_repeat_dataset( input_resource, buffer_size=self._buffer_size, count=self._count, - seed=self._seed, - seed2=self._seed2, + seed=seed, + seed2=seed2, **dataset_ops.flat_structure(self)) # pylint: enable=protected-access |