Skip to content

GitLab

Commit edcb2146 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal changes. 

PiperOrigin-RevId: 304633832
parent 4a998743
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
official/vision/detection/dataloader/input_reader.py
View file @ edcb2146
...@@ -85,13 +85,14 @@ class InputFn(object): ...@@ -85,13 +85,14 @@ class InputFn(object):
if self._input_sharding and ctx and ctx.num_input_pipelines > 1: if self._input_sharding and ctx and ctx.num_input_pipelines > 1:
dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id) dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
dataset = dataset.cache()
if self._is_training: if self._is_training:
dataset = dataset.repeat() dataset = dataset.repeat()
dataset = dataset.interleave( dataset = dataset.interleave(
map_func=lambda file_name: self._dataset_fn(file_name), cycle_length=32, map_func=lambda file_name: self._dataset_fn(file_name), cycle_length=32,
num_parallel_calls=tf.data.experimental.AUTOTUNE) num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.cache()
if self._is_training: if self._is_training:
# Large shuffle size is critical for 2vm input pipeline. Can use small # Large shuffle size is critical for 2vm input pipeline. Can use small
......
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