tensorflow通过tfrecord高效读写数据

利用tensorflow提供的tfrecord数据存储格式工具，我们可以将我们已经进行过处理的数据保存起来，以便我们下次更高效地读取，略过数据处理的过程，提高效率。具体的步骤大概分为以下几步：

1. 将数据转化为tf.train.Feature，然后存于字典；
2. 接着，将其转化为tf.train.example，然后进行序列化，写入tf.python_io.TFRecordWriter，到这里就完成了写入的操作；
3. 读取的时候，首先通过tf.data.TFRecordDataset来读取它；
4. 然后，通过tf.parse_single_example来解析还原数据原本的结构；
5. 最后就可以结合我们上次提高的batch来批量获取。（关于batch、shuffle、repeat函数的详细介绍TensorFlow dataset.shuffle、batch、repeat用法）

写入TFRecordWriter

import tensorflow as tf
import collections
import numpy as np

inputs_1 = np.array([
        [[1,2], [3,4]],
        [[5,6], [7,8]]
        ], dtype=np.int32)  # 这里的类型int32要与解析时使用的类型保持一致
inputs_2 = [
        [1.1, 2.2, 3.3], 
        [4.4, 5.5, 6.6]
        ]
lables = [0, 1]

def create_int_feature(values):
  f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))  # 需要注意这里接受的格式是list，并且只能是一维的
  return f

def create_float_feature(values):
  f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
  return f

def create_bytes_feature(values):
  f = tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))
  return f

writer = tf.python_io.TFRecordWriter('test.tfrecord')  # test.tfrecord是写入的文件路径

for i1, i2, l in zip(inputs_1, inputs_2, lables):
    features = collections.OrderedDict()  # 这里是新建一个有序字典
    # 对于多维数组，只能先将其转化为byte，才能传递给Feature
    features['inputs_1'] = create_bytes_feature(i1.tostring())
    features['inputs_2'] = create_float_feature(i2)
    features['labels'] = create_int_feature([l])
    
    example = tf.train.Example(features=tf.train.Features(feature=features))
    
    writer.write(example.SerializeToString())
writer.close()

读取并解析为dataset

源数据的数据类型最好与解析时使用的类型保持一致，不然可能会出现错误。
如inputs_1（ndarray）时使用的类型时np.int32，那么在tf.decode_raw进行解析时就相应地使用tf.int32。

name_to_features = {
      "inputs_1": tf.FixedLenFeature([], tf.string),  
      "inputs_2": tf.FixedLenFeature([3], tf.float32),  # 这里的格式需要与写入的保持一致，否则可能出现解析错误
      "labels": tf.FixedLenFeature([], tf.int64)
  }

d = tf.data.TFRecordDataset('test.tfrecord')
d = d.repeat() # 这里repeat不传递参数，则会无限重复
d = d.shuffle(buffer_size=2)
# map_and_batch其实就是map和batch结合在一起而已
d = d.apply(tf.contrib.data.map_and_batch(
        lambda record: tf.parse_single_example(record, name_to_features),
        batch_size=1))

iters = d.make_one_shot_iterator()
batch = iters.get_next()

# BytesList解析时会丢失shape信息，需要自己还原它的shape，所以一般也会将shape数据一同写入
# 这里需要将byte解析成原本的数据结构，这里的tf.int32需要与写入时的格式保持一致
inputs_1_batch = tf.decode_raw(batch['inputs_1'], tf.int32)  # tf.int32类型要与源数据的类型保持一致
inputs_1_batch = tf.reshape(inputs_1_batch, [-1, 2, 2])  
# 因为每次是batch个inputs_1，所以shape是[-1, 2, 2]，原来的shape是[2, 2]
inputs_2_batch = batch['inputs_2']
labels_batch = batch['labels']

sess = tf.Session()
# 这样我们就可以每次获取一个batch的数据了
sess.run([inputs_1_batch, inputs_2_batch, labels_batch])

欢迎关注同名公众号：“我就算饿死也不做程序员”。
交个朋友，一起交流，一起学习，一起进步。