Tensorflow数据读写：Numpy存储为TFRecord文件与读取

Tensorflow数据读写：Numpy存储为TFRecord文件与读取

用Tensorflow训练模型，读取数据有三种方法：

• 每个epoch/batch将内存中的Numpy数据送入placeholder，只适用于小数据集，会相当占用内存。
• 从硬盘里的txt或csv文件读取，IO操作比较耗时。
• 读取tensorflow推荐的TFRecord文件

  TFRecord文件是一种能将data和label一起存储的二进制文件，能更好地利用内存，在tensorflow的graph中更快地复制、移动、读取。TFRecord文件包含了tf.train.Example 协议缓冲区（protocol buffer），可以先将数据转成字符串序列化，填入到协议缓冲区，再由TFRecordWritier写入TFRecord文件。

---

Numpy存TFRecord

def save_tfrecords(data, label, desfile):
    with tf.python_io.TFRecordWriter(desfile) as writer:
        for i in range(len(data)):
            features = tf.train.Features(
                feature = {
                    "data":tf.train.Feature(bytes_list = tf.train.BytesList(value = [data[i].astype(np.float64).tostring()])),
                    "label":tf.train.Feature(int64_list = tf.train.Int64List(value = [label[i]]))
                }
            )
            example = tf.train.Example(features = features)
            serialized = example.SerializeToString()
            writer.write(serialized)

使用举例

例如，我们把一个含有10个样本，维度不一的数据集，跟它们的label一起存储到tfrecord文件中。

# 将不定长样本padding补0成定长
def padding(data, maxlen=10):
    for i in range(len(data)):
        data[i] = np.hstack([data[i], np.zeros((maxlen-len(data[i])))])

lens = np.random.randint(low=3,high=10,size=(10,))
data = [np.arange(l) for l in lens]
padding(data)
label = [0,0,0,0,0,1,1,1,1,1]

save_tfrecords(data, label, "./data.tfrecords")

---

读TFRecord转Numpy

def _parse_function(example_proto):
  features = {
  
  "data": tf.FixedLenFeature((), tf.string),
              "label": tf.FixedLenFeature((), tf.int64)}
  parsed_features = tf.parse_single_example(example_proto, features)
  data = tf.decode_raw(parsed_features['data'], tf.float32)
  return data, parsed_features["label"]

def load_tfrecords(srcfile):
    sess = tf.Session()

    dataset = tf.data.TFRecordDataset(srcfile) # load tfrecord file
    dataset = dataset.map(_parse_function) # parse data into tensor
    dataset = dataset.repeat(2) # repeat for 2 e