Merge pull request #2080 from mari-linhares/patch-5

Refactoring and making just one tf.record file for training for efficiency

tutorials/image/cifar10_estimator/cifar10.py

@@ -39,15 +39,8 @@ class Cifar10DataSet(object):
     self.use_distortion = use_distortion
 
   def get_filenames(self):
-    if self.subset == 'train':
-      return [
-          os.path.join(self.data_dir, 'data_batch_%d.tfrecords' % i)
-          for i in xrange(1, 5)
-      ]
-    elif self.subset == 'validation':
-      return [os.path.join(self.data_dir, 'data_batch_5.tfrecords')]
-    elif self.subset == 'eval':
-      return [os.path.join(self.data_dir, 'test_batch.tfrecords')]
+    if self.subset in ['train', 'validation', 'eval']:
+      return [os.path.join(self.data_dir, self.subset + '.tfrecords')]
     else:
       raise ValueError('Invalid data subset "%s"' % self.subset)
 
@@ -66,7 +59,9 @@ class Cifar10DataSet(object):
     image.set_shape([DEPTH * HEIGHT * WIDTH])
 
     # Reshape from [depth * height * width] to [depth, height, width].
-    image = tf.transpose(tf.reshape(image, [DEPTH, HEIGHT, WIDTH]), [1, 2, 0])
+    image = tf.cast(
+        tf.transpose(tf.reshape(image, [DEPTH, HEIGHT, WIDTH]), [1, 2, 0]),
+        tf.float32)
     label = tf.cast(features['label'], tf.int32)
 
     # Custom preprocessing .

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,18 +61,12 @@ def read_pickle_from_file(filename):
   return data_dict
 
 
-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')
-
-    print('Generating %s' % output_file)
-
-    record_writer = tf.python_io.TFRecordWriter(output_file)
+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']
@@ -82,8 +79,18 @@ def main(argv):
               'label': _int64_feature(labels[i])
           }))
       record_writer.write(example.SerializeToString())
-    record_writer.close()
+  record_writer.close()
 
+def main(argv):
+  del argv  # Unused.
+
+  file_names = _get_file_names()
+  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_files, output_file)
   print('Done!')