Generating a single file for each mode: train, validation, eval

tutorials/image/cifar10_estimator/generate_cifar10_tfrecords.py

@@ -46,9 +46,12 @@ def _bytes_feature(value):
 
 
 def _get_file_names():
-  """Returns the file names expected to exist in the input_dir."""
-  file_names = ['data_batch_%d' % i for i in xrange(1, 6)]
-  file_names.append('test_batch')
+  """Returns the file names expected to exist for training, validation
+  and evaluation in the input_dir."""
+  file_names = {}
+  file_names['train'] = ['data_batch_%d' % i for i in xrange(1, 5)]
+  file_names['validation'] = ['data_batch_5']
+  file_names['eval'] = ['test_batch']
   return file_names
 
 
@@ -58,11 +61,12 @@ def read_pickle_from_file(filename):
   return data_dict
 
 
-def convert_to_tfrecord(input_file, name):
+def convert_to_tfrecord(input_files, output_file):
   """Converts a file to tfrecords."""
   print('Generating %s' % output_file)
   record_writer = tf.python_io.TFRecordWriter(output_file)
 
+  for input_file in input_files:
     data_dict = read_pickle_from_file(input_file)
     data = data_dict['data']
     labels = data_dict['labels']
@@ -81,12 +85,12 @@ def main(argv):
   del argv  # Unused.
 
   file_names = _get_file_names()
-  for file_name in file_names:
-    input_file = os.path.join(FLAGS.input_dir, file_name)
-    output_file = os.path.join(FLAGS.output_dir, file_name + '.tfrecords')
+  for mode, files in file_names.items():
+    input_files = [
+        os.path.join(FLAGS.input_dir, f) for f in files]
+    output_file = os.path.join(FLAGS.output_dir, mode + '.tfrecords')
     # Convert to Examples and write the result to TFRecords.
-    convert_to_tfrecord(input_file, output_file)
-
+    convert_to_tfrecord(input_files, output_file)
   print('Done!')