"examples/vscode:/vscode.git/clone" did not exist on "26e80e014331410087a76cc7979ea99fb736f30a"
Commit ed877c00 authored by Zhichao Lu's avatar Zhichao Lu Committed by pkulzc
Browse files

Expose read block length via input reader proto so it can be tuned.

PiperOrigin-RevId: 190292280
parent 2a43dca3
...@@ -51,6 +51,9 @@ message InputReader { ...@@ -51,6 +51,9 @@ message InputReader {
// Number of reader instances to create. // Number of reader instances to create.
optional uint32 num_readers = 6 [default=32]; optional uint32 num_readers = 6 [default=32];
// Number of records to read from each reader at once.
optional uint32 read_block_length = 15 [default=32];
// Number of decoded records to prefetch before batching. // Number of decoded records to prefetch before batching.
optional uint32 prefetch_size = 13 [default = 512]; optional uint32 prefetch_size = 13 [default = 512];
......
...@@ -127,7 +127,8 @@ def read_dataset(file_read_func, decode_func, input_files, config): ...@@ -127,7 +127,8 @@ def read_dataset(file_read_func, decode_func, input_files, config):
records_dataset = filename_dataset.apply( records_dataset = filename_dataset.apply(
tf.contrib.data.parallel_interleave( tf.contrib.data.parallel_interleave(
file_read_func, cycle_length=config.num_readers, sloppy=True)) file_read_func, cycle_length=config.num_readers,
block_length=config.read_block_length, sloppy=True))
if config.shuffle: if config.shuffle:
records_dataset.shuffle(config.shuffle_buffer_size) records_dataset.shuffle(config.shuffle_buffer_size)
tensor_dataset = records_dataset.map( tensor_dataset = records_dataset.map(
......
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