tf.data.Dataset读取TFRecord文件进行训练

前言

之前跑了一个小的CNN，因为训练数据量也不大，自己的电脑显卡2g，直接将数据读到内存中再取batch来训练还能直接跑；后来数据集大了一些，再用这种方法就不行了，内存不足。没办法，找了一些TFRecord和Dataset的资料把程序改了改，这样就可以边取数据边训练，总算能跑了，以后再也不用怕训练数据多啦。

数据

我的训练数据是一些元素为float类型的矩阵，用numpy生成一些来代表吧：

data = np.random.rand(5,10,10)

这样就生成了shape为（5,10,10），元素类型为float的矩阵data了。意义为5个大小为（10,10）的矩阵。

生成TFRecord文件

直接上代码：

import numpy as np
import tensorflow as tf

data = np.random.rand(5,10,10)

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

# 将数据转化为tf.train.Example格式。
def _make_example(data_raw):
    data_raw = data.tostring()# 先将元素为float类型的numpy数组转换为string
    example = tf.train.Example(features=tf.train.Features(feature={
        'data_raw': _bytes_feature(data_raw)
    }))
    return example

num_examples = data.shape[0]

# 输出包含训练数据的TFRecord文件。
with tf.python_io.TFRecordWriter("output.tfrecords") as writer:
    for index in range(num_examples):
        example = _make_example(data[index])
        writer.write(example.SerializeToString())
print("TFRecord训练文件已保存。")

运行完后就会在当前目录下生成output.tfrecords的文件了，这就是我们的数据。虽然感觉不好理解，其实也就这几句代码就完成了TFRecord数据的生成。
因为TFRecord只能存储BytesList、FloatList、Int64List，所以矩阵不能直接存储，可以将矩阵转换成字符串来存储，如以上程序data_raw = data.tostring()所示。也可以将矩阵flat成一维即FloatList来存储，读出时再reshape转换为之前的格式，按理可行，但我没试。

使用Dataset恢复数据

def parser(record1):
	features = tf.parse_single_example(
    record1,
    features={
        'data_raw':tf.FixedLenFeature([],tf.string),
    })
	data_raw = features['data_raw']
	return data_raw

batch_size = 2
dataset = tf.data.TFRecordDataset("output.tfrecords")
dataset = dataset.map(parser)
dataset = dataset.batch(batch_size)
NUM_EPOCHS = 2
dataset = dataset.repeat(NUM_EPOCHS)
iterator = dataset.make_initializable_iterator()
next_element = iterator.get_next()

with tf.Session() as sess:
    x_batch = []
    sess.run((tf.global_variables_initializer(),
			tf.local_variables_initializer()))
    sess.run(iterator.initializer)
    while True:
        try:
            x_train = sess.run(next_element) # 取出设定的
            for i in range(x_train.shape[0]):
                x_batch.append(list(np.fromstring(x_train[i],np.float32)))
            print(x_batch)
        except tf.errors.OutOfRangeError:
            break

x_batch就是恢复出来的数据了，是list列表，根据需要可以转换成其他的形式。
parser解析函数的作用是把数据从TFRecord文件中解析出来。
map（parser）表示对数据集中的每一条数据调用parser方法。
dataset.batch函数是设置一次取出的batch的大小。
dataset.repeat函数表示取全部数据几个epoch，比如我这里设置为2，那么在取完两次全部数据后，就会抛出OutOfRangeError异常，可以以此来停止程序循环。
dataset.shuffle函数将数据集打乱顺序，这里没有展示。
每运行一次sess.run(next_element)就会从TFRecord文件中取出设定的batch大小的数据。取出来的数据的格式是字符串，通过np.fromstring()函数转换成真正的数据。这里要注意转换成的float类型，你原来的数据是float32你转换时就要转换为float32，你如果转换成float64，那肯定数据转换出来是错的。

未理解的地方

取出一个batch的数据之后，如果直接整体操作：

x_train = np.fromstring(x_train,np.float32)

转换出来的数据是不对的，必须要一条一条的转换出来，然后再组合起来才行。