Merge pull request tensorflow#2732 from tensorflow/wide-deep-shuffle-…

…size Shuffle an exact epoch of 32561 examples in wide_deep.py rather than 100k
rockokw · Nov 8, 2017 · a6e1b89 · a6e1b89
2 parents 5f0776a + 6d44720
commit a6e1b89
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/official/wide_deep/wide_deep.py b/official/wide_deep/wide_deep.py
@@ -61,7 +61,10 @@
     '--test_data', type=str, default='/tmp/census_data/adult.test',
     help='Path to the test data.')
 
-_SHUFFLE_BUFFER = 100000
+_NUM_EXAMPLES = {
+    'train': 32561,
+    'validation': 16281,
+}
 
 
 def build_model_columns():
@@ -181,7 +184,7 @@ def parse_csv(value):
   dataset = tf.data.TextLineDataset(data_file)
 
   if shuffle:
-    dataset = dataset.shuffle(buffer_size=_SHUFFLE_BUFFER)
+    dataset = dataset.shuffle(buffer_size=_NUM_EXAMPLES['train'])
 
   dataset = dataset.map(parse_csv, num_parallel_calls=5)