Skip to content

GitLab

Commit ada2ed77 authored by Yeqing Li's avatar Yeqing Li Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 277960691
parent 37c12026
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 6 deletions
official/vision/detection/dataloader/input_reader.py
View file @ ada2ed77
...@@ -79,11 +79,9 @@ class InputFn(object): ...@@ -79,11 +79,9 @@ class InputFn(object):
if self._is_training: if self._is_training:
dataset = dataset.repeat() dataset = dataset.repeat()
dataset = dataset.apply( dataset = dataset.interleave(
tf.data.experimental.parallel_interleave( map_func=lambda file_name: self._dataset_fn(file_name), cycle_length=32,
lambda file_name: self._dataset_fn(file_name).prefetch(1), num_parallel_calls=tf.data.experimental.AUTOTUNE)
cycle_length=32,
sloppy=self._is_training))
if self._is_training: if self._is_training:
dataset = dataset.shuffle(64) dataset = dataset.shuffle(64)
...@@ -91,7 +89,8 @@ class InputFn(object): ...@@ -91,7 +89,8 @@ class InputFn(object):
dataset = dataset.take(self._num_examples) dataset = dataset.take(self._num_examples)
# Parses the fetched records to input tensors for model function. # Parses the fetched records to input tensors for model function.
dataset = dataset.map(self._parser_fn, num_parallel_calls=64) dataset = dataset.map(
self._parser_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.batch(batch_size, drop_remainder=True) dataset = dataset.batch(batch_size, drop_remainder=True)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
return dataset return dataset
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment