[Modelzoo] Change prefetch setting in data processing. (DeepRec-AI#212)

Signed-off-by: Duyi-Wang <[email protected]>
StevenShi-23 · May 11, 2022 · 61020d2 · 61020d2
1 parent 08fc710
commit 61020d2
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 12 deletions.
diff --git a/modelzoo/DIEN/train.py b/modelzoo/DIEN/train.py
@@ -640,11 +640,10 @@ def parse_csv(value, neg_value):
     dataset = dataset.shuffle(buffer_size=20000,
                               seed=args.seed)  # set seed for reproducing
     dataset = dataset.repeat(num_epochs)
-    dataset = dataset.prefetch(batch_size)
     dataset = dataset.batch(batch_size)
     dataset = dataset.map(parse_csv,
                           num_parallel_calls=tf.data.experimental.AUTOTUNE)
-    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(2)
     return dataset
 
 

diff --git a/modelzoo/DIN/train.py b/modelzoo/DIN/train.py
@@ -467,11 +467,10 @@ def parse_csv(value):
     dataset = dataset.shuffle(buffer_size=20000,
                               seed=args.seed)  # set seed for reproducing
     dataset = dataset.repeat(num_epochs)
-    dataset = dataset.prefetch(batch_size)
     dataset = dataset.batch(batch_size)
     dataset = dataset.map(parse_csv,
                           num_parallel_calls=tf.data.experimental.AUTOTUNE)
-    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(2)
     return dataset
 
 

diff --git a/modelzoo/DLRM/train.py b/modelzoo/DLRM/train.py
@@ -315,10 +315,9 @@ def parse_csv(value):
     dataset = dataset.shuffle(buffer_size=20000,
                               seed=args.seed)  # set seed for reproducing
     dataset = dataset.repeat(num_epochs)
-    dataset = dataset.prefetch(batch_size)
     dataset = dataset.batch(batch_size)
     dataset = dataset.map(parse_csv, num_parallel_calls=28)
-    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(2)
     return dataset
 
 

diff --git a/modelzoo/DSSM/train.py b/modelzoo/DSSM/train.py
@@ -291,10 +291,9 @@ def parse_csv(value):
     dataset = dataset.shuffle(buffer_size=20000,
                               seed=args.seed)  # fix seed for reproducing
     dataset = dataset.repeat(num_epochs)
-    dataset = dataset.prefetch(batch_size)
     dataset = dataset.batch(batch_size)
     dataset = dataset.map(parse_csv, num_parallel_calls=28)
-    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(2)
     return dataset
 
 

diff --git a/modelzoo/DeepFM/train.py b/modelzoo/DeepFM/train.py
@@ -275,10 +275,9 @@ def parse_csv(value):
     dataset = dataset.shuffle(buffer_size=20000,
                               seed=args.seed)  # fix seed for reproducing
     dataset = dataset.repeat(num_epochs)
-    dataset = dataset.prefetch(batch_size)
     dataset = dataset.batch(batch_size)
     dataset = dataset.map(parse_csv, num_parallel_calls=28)
-    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(2)
     return dataset
 
 

diff --git a/modelzoo/WDL/train.py b/modelzoo/WDL/train.py
@@ -318,10 +318,9 @@ def parse_csv(value):
     dataset = dataset.shuffle(buffer_size=20000,
                               seed=args.seed)  # fix seed for reproducing
     dataset = dataset.repeat(num_epochs)
-    dataset = dataset.prefetch(batch_size)
     dataset = dataset.batch(batch_size)
     dataset = dataset.map(parse_csv, num_parallel_calls=28)
-    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(2)
     return dataset