Update train.py

tensorlayer · Jul 25, 2019 · bc02a02 · bc02a02
1 parent 849ab78
commit bc02a02
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -65,7 +65,7 @@ def _map_fn_train(img):
     train_ds = train_ds.map(_map_fn_train, num_parallel_calls=multiprocessing.cpu_count())
     train_ds = train_ds.repeat(n_epoch_init + n_epoch)
     train_ds = train_ds.shuffle(shuffle_buffer_size)
-    train_ds = train_ds.prefetch(buffer_size=4096)
+    train_ds = train_ds.prefetch(buffer_size=2)
     train_ds = train_ds.batch(batch_size)
     # value = train_ds.make_one_shot_iterator().get_next()