A `Dataset` that caches elements of its input.
| 3084 | |
| 3085 | |
| 3086 | class CacheDataset(UnaryUnchangedStructureDataset): |
| 3087 | """A `Dataset` that caches elements of its input.""" |
| 3088 | |
| 3089 | def __init__(self, input_dataset, filename): |
| 3090 | """See `Dataset.cache()` for details.""" |
| 3091 | self._input_dataset = input_dataset |
| 3092 | self._filename = ops.convert_to_tensor( |
| 3093 | filename, dtype=dtypes.string, name="filename") |
| 3094 | if tf2.enabled() and (context.executing_eagerly() or |
| 3095 | ops.get_default_graph()._building_function): # pylint: disable=protected-access |
| 3096 | self._cache = _MemoryCache() |
| 3097 | variant_tensor = gen_dataset_ops.cache_dataset_v2( |
| 3098 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3099 | filename=self._filename, |
| 3100 | cache=self._cache.handle, |
| 3101 | **self._flat_structure) |
| 3102 | else: |
| 3103 | variant_tensor = gen_dataset_ops.cache_dataset( |
| 3104 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3105 | filename=self._filename, |
| 3106 | **self._flat_structure) |
| 3107 | super(CacheDataset, self).__init__(input_dataset, variant_tensor) |
| 3108 | |
| 3109 | |
| 3110 | class _RandomSeedGeneratorDeleter(object): |