Commit 8a5702f2 authored by Ruoxin Sang's avatar Ruoxin Sang Committed by A. Unique TensorFlower
Browse files

Internal change

PiperOrigin-RevId: 276103495
parent 444e8c79
...@@ -104,21 +104,12 @@ def process_record_dataset(dataset, ...@@ -104,21 +104,12 @@ def process_record_dataset(dataset,
logging.info( logging.info(
'datasets_num_private_threads: %s', datasets_num_private_threads) 'datasets_num_private_threads: %s', datasets_num_private_threads)
# Disable intra-op parallelism to optimize for throughput instead of latency.
options = tf.data.Options()
options.experimental_threading.max_intra_op_parallelism = 1
dataset = dataset.with_options(options)
# Prefetches a batch at a time to smooth out the time taken to load input
# files for shuffling and processing.
dataset = dataset.prefetch(buffer_size=batch_size)
if is_training: if is_training:
# Shuffles records before repeating to respect epoch boundaries. # Shuffles records before repeating to respect epoch boundaries.
dataset = dataset.shuffle(buffer_size=shuffle_buffer) dataset = dataset.shuffle(buffer_size=shuffle_buffer)
# Repeats the dataset for the number of epochs to train. # Repeats the dataset for the number of epochs to train.
dataset = dataset.repeat() dataset = dataset.repeat()
# Parses the raw records into images and labels. # Parses the raw records into images and labels.
dataset = dataset.map( dataset = dataset.map(
lambda value: parse_record_fn(value, is_training, dtype), lambda value: parse_record_fn(value, is_training, dtype),
......
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