A `Dataset` that asynchronously prefetches its input.
| 3735 | |
| 3736 | |
| 3737 | class PrefetchDataset(UnaryUnchangedStructureDataset): |
| 3738 | """A `Dataset` that asynchronously prefetches its input.""" |
| 3739 | |
| 3740 | def __init__(self, input_dataset, buffer_size, slack_period=None): |
| 3741 | """See `Dataset.prefetch()` for details. |
| 3742 | |
| 3743 | Args: |
| 3744 | input_dataset: The input dataset. |
| 3745 | buffer_size: See `Dataset.prefetch()` for details. |
| 3746 | slack_period: (Optional.) An integer. If non-zero, determines the number |
| 3747 | of GetNext calls before injecting slack into the execution. This may |
| 3748 | reduce CPU contention at the start of a step. Note that a tensorflow |
| 3749 | user should not have to set this manually; enable this behavior |
| 3750 | automatically via `tf.data.Options.experimental_slack` instead. Defaults |
| 3751 | to None. |
| 3752 | """ |
| 3753 | self._input_dataset = input_dataset |
| 3754 | if buffer_size is None: |
| 3755 | buffer_size = -1 # This is the sentinel for auto-tuning. |
| 3756 | self._buffer_size = ops.convert_to_tensor( |
| 3757 | buffer_size, dtype=dtypes.int64, name="buffer_size") |
| 3758 | |
| 3759 | with ops.device(input_dataset._variant_tensor.device): |
| 3760 | variant_tensor = gen_dataset_ops.prefetch_dataset( |
| 3761 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3762 | buffer_size=self._buffer_size, |
| 3763 | slack_period=slack_period, |
| 3764 | **self._flat_structure) |
| 3765 | |
| 3766 | super(PrefetchDataset, self).__init__(input_dataset, variant_tensor) |
| 3767 | |
| 3768 | |
| 3769 | class WindowDataset(UnaryDataset): |