experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Re-prefetch the dataset for V1Adapater

Browse Source 
PiperOrigin-RevId: 226650870
This commit is contained in:
Jianwei Xie 2018-12-22 21:40:58 -08:00 committed by TensorFlower Gardener
parent 3785cffbdb
commit 0a708f8da4
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/python/distribute/values.py
View File

@ -1959,6 +1959,9 @@ def _split_dataset_batch(dataset, split_batch_by):
prefetch_buffer = None prefetch_buffer = None
if isinstance(dataset, dataset_ops.PrefetchDataset): if isinstance(dataset, dataset_ops.PrefetchDataset):
prefetch_buffer = dataset._buffer_size prefetch_buffer = dataset._buffer_size
elif (isinstance(dataset, dataset_ops.DatasetV1Adapter)
and isinstance(dataset._dataset, dataset_ops.PrefetchDataset)):
prefetch_buffer = dataset._dataset._buffer_size
# pylint: enable=protected-access # pylint: enable=protected-access
if tensor_util.is_tensor(batch_size): if tensor_util.is_tensor(batch_size):