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,35 +61,36 @@ 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)
 
-  data_dict = read_pickle_from_file(input_file)
-  data = data_dict['data']
-  labels = data_dict['labels']
-
-  num_entries_in_batch = len(labels)
-  for i in range(num_entries_in_batch):
-    example = tf.train.Example(
-        features=tf.train.Features(feature={
-            'image': _bytes_feature(data[i].tobytes()),
-            'label': _int64_feature(labels[i])
-        }))
-    record_writer.write(example.SerializeToString())
+  for input_file in input_files:
+    data_dict = read_pickle_from_file(input_file)
+    data = data_dict['data']
+    labels = data_dict['labels']
+
+    num_entries_in_batch = len(labels)
+    for i in range(num_entries_in_batch):
+      example = tf.train.Example(
+          features=tf.train.Features(feature={
+              'image': _bytes_feature(data[i].tobytes()),
+              'label': _int64_feature(labels[i])
+          }))
+      record_writer.write(example.SerializeToString())
   record_writer.close()
 
 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!')