1# Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Experimental shuffle ops.""" 16from tensorflow.python.data.ops import dataset_ops 17from tensorflow.python.data.util import random_seed 18from tensorflow.python.framework import constant_op 19from tensorflow.python.framework import dtypes 20from tensorflow.python.framework import ops 21from tensorflow.python.ops import gen_dataset_ops 22from tensorflow.python.util import deprecation 23from tensorflow.python.util.tf_export import tf_export 24 25 26class _ShuffleAndRepeatDataset(dataset_ops.UnaryUnchangedStructureDataset): 27 """A `Dataset` that fuses `shuffle` and `repeat`.""" 28 29 def __init__(self, input_dataset, buffer_size, count=None, seed=None): 30 self._input_dataset = input_dataset 31 self._buffer_size = ops.convert_to_tensor( 32 buffer_size, dtype=dtypes.int64, name="buffer_size") 33 if count is None: 34 self._count = constant_op.constant(-1, dtype=dtypes.int64, name="count") 35 else: 36 self._count = ops.convert_to_tensor( 37 count, dtype=dtypes.int64, name="count") 38 self._seed, self._seed2 = random_seed.get_seed(seed) 39 variant_tensor = gen_dataset_ops.shuffle_and_repeat_dataset( 40 self._input_dataset._variant_tensor, # pylint: disable=protected-access 41 buffer_size=self._buffer_size, 42 count=self._count, 43 seed=self._seed, 44 seed2=self._seed2, 45 **self._flat_structure) 46 super(_ShuffleAndRepeatDataset, self).__init__(input_dataset, 47 variant_tensor) 48 49 50@deprecation.deprecated( 51 None, 52 "Use `tf.data.Dataset.shuffle(buffer_size, seed)` followed by " 53 "`tf.data.Dataset.repeat(count)`. Static tf.data optimizations will take " 54 "care of using the fused implementation.") 55@tf_export("data.experimental.shuffle_and_repeat") 56def shuffle_and_repeat(buffer_size, count=None, seed=None): 57 """Shuffles and repeats a Dataset, reshuffling with each repetition. 58 59 >>> d = tf.data.Dataset.from_tensor_slices([1, 2, 3]) 60 >>> d = d.apply(tf.data.experimental.shuffle_and_repeat(2, count=2)) 61 >>> [elem.numpy() for elem in d] # doctest: +SKIP 62 [2, 3, 1, 1, 3, 2] 63 64 ```python 65 dataset.apply( 66 tf.data.experimental.shuffle_and_repeat(buffer_size, count, seed)) 67 ``` 68 69 produces the same output as 70 71 ```python 72 dataset.shuffle( 73 buffer_size, seed=seed, reshuffle_each_iteration=True).repeat(count) 74 ``` 75 76 In each repetition, this dataset fills a buffer with `buffer_size` elements, 77 then randomly samples elements from this buffer, replacing the selected 78 elements with new elements. For perfect shuffling, set the buffer size equal 79 to the full size of the dataset. 80 81 For instance, if your dataset contains 10,000 elements but `buffer_size` is 82 set to 1,000, then `shuffle` will initially select a random element from 83 only the first 1,000 elements in the buffer. Once an element is selected, 84 its space in the buffer is replaced by the next (i.e. 1,001-st) element, 85 maintaining the 1,000 element buffer. 86 87 Args: 88 buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the maximum 89 number elements that will be buffered when prefetching. 90 count: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the number 91 of times the dataset should be repeated. The default behavior (if `count` 92 is `None` or `-1`) is for the dataset be repeated indefinitely. 93 seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the random 94 seed that will be used to create the distribution. See 95 `tf.random.set_seed` for behavior. 96 97 Returns: 98 A `Dataset` transformation function, which can be passed to 99 `tf.data.Dataset.apply`. 100 """ 101 102 def _apply_fn(dataset): # pylint: disable=missing-docstring 103 return _ShuffleAndRepeatDataset(dataset, buffer_size, count, seed) 104 105 return _apply_fn 106