使用TensorFlow Dataset读取数据

在使用TensorFlow构建模型并进行训练时，如何读取数据并将数据恰当地送进模型，是一个首先需要考虑的问题。以往通常所用的方法无外乎以下几种：

1.建立placeholder，然后使用feed_dict将数据feed进placeholder进行使用。使用这种方法十分灵活，可以一下子将所有数据读入内存，然后分batch进行feed；也可以建立一个Python的generator，一个batch一个batch的将数据读入，并将其feed进placeholder。这种方法很直观，用起来也比较方便灵活，但是这种方法的效率较低，难以满足高速计算的需求。

2.使用TensorFlow的QueueRunner，通过一系列的Tensor操作，将磁盘上的数据分批次读入并送入模型进行使用。这种方法效率很高，但因为其牵涉到Tensor操作，不够直观，也不方便调试，所有有时候会显得比较困难。使用这种方法时，常用的一些操作包括tf.TextLineReader，tf.FixedLengthRecordReader以及tf.decode_raw等等。如果需要循环，条件操作，还需要使用TensorFlow的tf.while_loop，tf.case等操作，更是难上加难。

因此，在这种情况下，TensorFlow在后续的更新中，自1.x版本开始，逐步开发引入了tf.data.Dataset模块，使其数据读入的操作变得更为方便，而支持多线程（进程）的操作，也在效率上获得了一定程度的提高。本文就将使用tf.data.Dataset过程中的一些经验进行总结记录，以便备忘。

如我们所知，在使用TensorFlow建立模型进行训练的时候，可以很容易生成这样的文件，来表示数据：

1.  data/01.jpg,猫
2.  data/05.jpg,狗
3.  data/03.jpg,猫
4.  data/04.jpg,狗
5.  data/06.jpg,狗
6.  data/02.jpg,猫

这种数据格式可以很方便地进行各种操作，比如划分数据集、shuffle等等。所以我们就以将这样的数据通过tf.data.Dataset读入进行训练为例，来讲述其用法。

具体来说，使用tf.data.Dataset读取数据，本文讲述这样三种方法：

1.首先将数据读入内存，然后使用tf.data.Dataset构建数据集

具体来说，因为tf.data.Dataset.from_tensor_slices()函数会对tensor和numpy array的处理一视同仁，所以该函数既可以使用tensor参数，也可以直接使用numpy array作参数，使用numpy array作参数，即是第1种方法。

如下所示：

1.  images = ...
2.  labels = ...
3.  data = tf.data.Dataset.from_tensor_slices((images, labels))
4.  data = data.batch(batch_size)
5.  iterator = tf.data.Iterator.from_structure(data.output_types,
6.  data.output_shapes)
7.  init_op = iterator.make_initializer(data)
8.  with tf.Session()  as sess:
9.  sess.run(init_op)
10.  try:
11.  images, labels = iterator.get_next()
12.  except tf.errors.OutOfRangeError:
13.  sess.run(init_op)

第1~2行，首先，将数据images、labels读入内存；

第3~4行，使用读入内存的数据images、labels构建Dataset，并设置Dataset的batch大小；

第5行，基于此前构建的Dataset的数据类型和结构，构建一个iterator；

第6行，基于此前构建的Dataset构建一个初始化op。

随后的操作，即是在TensorFlow的session里，首先进行初始化操作，然后即可通过iterator的函数逐批获得数据，并进行使用了。

需要注意的是，iterator中的元素取完之后，会抛出OutOfRangeError异常，TensorFlow没有对这个异常进行处理，我们需要对其进行捕捉和处理。

本方法详细代码可参阅这里。

2.使用tf.data.Dataset包装一个generator读入数据

1中方法虽然简单，但其将数据一次读入，在面对大数据集时会束手无策。因此，我们可以建立一个读入数据的generator，然后使用tf.data.Dataset对其进行包装转换，即可实现逐batch读入数据的目的。如下：

1.  def gen():
2.  with  open('train.csv')  as f:
3.  lines = [line.strip().split(',')  for line in f.readlines()]
4.  index = 0
5.  while  True:
6.  image = cv2.imread(lines[index][0])
7.  image = cv2.resize(image, (224, 224))
8.  label = lines[index][1]
9.  yield  (image, label)
10.  index += 1
11.  if index == len(lines):
12.  index = 0

15.  batch_size = 2
16.  data = tf.data.Dataset.from_generator(gen, (tf.float32, tf.int32),
17.  (tf.TensorShape([224, 224, 3]), tf.TensorShape([])))
18.  data = data.batch(batch_size)
19.  iter = data.make_one_shot_iterator()
20.  with tf.Session()  as sess:
21.  images, labels = iter.get_next()

如上，首先构建一个generator：gen，然后使用tf.data.Dataset的from_generator函数，通过指定数据类型，数据的shape等参数，构建一个Dataset，当然，随后也要指定一下batch_size，最后使用make_one_shot_iterator()函数，构建一个iterator。

然后其使用方法即与前述相同了，不过需要说明的是，这里是通过一个永无尽头的generator构建的Dataset，所以其可以一直取数据，而不会出现1中所述的OutOfRange的问题。

本方法详细代码可参阅这里。

3.基于Tensor操作构建Dataset

前述两种方法，1中需要将数据一次全部读入内存，2中使用generator逐batch读入数据，虽然内存占用得到了控制，但是其效率仍然不高，读取速度较慢。在第3种方法里，我们通过TensorFlow提供的tensor操作来读取数据，并基于此，构建Dataset。

示例的代码片段如下：

1.  def _parse_function(filename, label):
2.  image_string = tf.read_file(filename)
3.  image_decoded = tf.image.decode_jpeg(image_string, channels=3)
4.  image = tf.cast(image_decoded, tf.float32)
5.  image = tf.image.resize_images(image, [224, 224])
6.  return image, filename, label

8.  images = tf.constant(image_names)
9.  labels = tf.constant(labels)
10.  images = tf.random_shuffle(images, seed=0)
11.  labels = tf.random_shuffle(labels, seed=0)
12.  data = tf.data.Dataset.from_tensor_slices((images, labels))

14.  data = data.map(_parse_function, num_parallel_calls=4)
15.  data = data.prefetch(buffer_size=batch_size * 10)
16.  data = data.batch(batch_size)

18.  iterator = tf.data.Iterator.from_structure(data.output_types,
19.  data.output_shapes)

21.  init_op = iterator.make_initializer(data)
22.  with tf.Session()  as sess:
23.  sess.run(init_op)
24.  try:
25.  images, filenames, labels = iterator.get_next()
26.  except tf.errors.OutOfRangeError:
27.  sess.run(init_op)

首先读入image names以及相应的labels，然后通过tf.constant构建constant Tensor：images, labels，并可选择地对其进行shuffle。

接着使用tf.data.Dataset.from_tensor_slices()函数基于images和labels构建Dataset。

然后使用map函数将函数应用到该Dataset上，本例中，将解析图像的函数_parse_function应用到Dataset上，还指定了多线程并行操作的线程数。

随后指定prefetch的buffer_size，以及batch的大小。

最后，基于构建的Dataset建立iterator，并定义iterator的初始化操作op，然后就可以按照正常的方式进行使用了。

需要注意的是，本方法构建的Dataset也会有OutOfRange的异常出现，需要恰当地进行捕捉并处理。

本方法详细代码可参阅这里。