Increase Prefetch in Example (#832)

* support channel_first * YAPF Corrections * yapf * change prefetech number * prefetch
tensorlayer · Sep 9, 2018 · 2014577 · 2014577
1 parent d42ef66
commit 2014577
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/examples/basic_tutorials/tutorial_cifar10_datasetapi.py b/examples/basic_tutorials/tutorial_cifar10_datasetapi.py
@@ -105,16 +105,16 @@ def _map_fn_test(img, target):
 ds = ds.map(_map_fn_train, num_parallel_calls=multiprocessing.cpu_count())
 ds = ds.repeat(n_epoch)
 ds = ds.shuffle(shuffle_buffer_size)
-ds = ds.prefetch(buffer_size=2)
+ds = ds.prefetch(buffer_size=4096)
 ds = ds.batch(batch_size)
 value = ds.make_one_shot_iterator().get_next()
 
 ds = tf.data.Dataset().from_generator(generator_test, output_types=(tf.float32,
                                                                     tf.int32))  # , output_shapes=((24, 24, 3), (1)))
+ds = ds.shuffle(shuffle_buffer_size)
 ds = ds.map(_map_fn_test, num_parallel_calls=multiprocessing.cpu_count())
 ds = ds.repeat(n_epoch)
-ds = ds.shuffle(shuffle_buffer_size)
-ds = ds.prefetch(buffer_size=2)
+ds = ds.prefetch(buffer_size=4096)
 ds = ds.batch(batch_size)
 value_test = ds.make_one_shot_iterator().get_next()
 

diff --git a/examples/data_process/tutorial_tf_dataset_voc.py b/examples/data_process/tutorial_tf_dataset_voc.py
@@ -81,11 +81,11 @@ def _map_fn(filename, annotation):
     image, annotation = tf.py_func(_data_aug_fn, [image, annotation], [tf.float32, tf.string])
     return image, annotation
 
-
 ds = tf.data.Dataset().from_generator(generator, output_types=(tf.string, tf.string))
+ds = ds.shuffle(shuffle_buffer_size)
 ds = ds.map(_map_fn, num_parallel_calls=multiprocessing.cpu_count())
 ds = ds.repeat(n_epoch)
-ds = ds.shuffle(shuffle_buffer_size)
+ds = ds.prefetch(buffer_size=2048)
 ds = ds.batch(batch_size)
 value = ds.make_one_shot_iterator().get_next()