tfrecord_model_test_mnist.py 文件源码

def read_and_decode(filename, one_hot=True, n_classes=None):
  """ Return tensor to read from TFRecord """
  filename_queue = tf.train.string_input_producer([filename])
  reader = tf.TFRecordReader()
  _, serialized_example = reader.read(filename_queue)
  features = tf.parse_single_example(serialized_example,
                                     features={
                                       'label': tf.FixedLenFeature([],
                                                                   tf.int64),
                                       'image_raw': tf.FixedLenFeature([],
                                                                       tf.string),
                                     })
  # You can do more image distortion here for training data
  img = tf.decode_raw(features['image_raw'], tf.uint8)
  img.set_shape([28 * 28])
  img = tf.reshape(img, [28, 28, 1])
  img = tf.cast(img, tf.float32) * (1. / 255) - 0.5

  label = tf.cast(features['label'], tf.int32)
  if one_hot and n_classes:
    label = tf.one_hot(label, n_classes)

  return img, label