使用TensorFlow Dataset读取数据备忘

在 2018-01-30 发布于 人工智能 下以来已有249人读过本文 | 0条评论 发表评论

在使用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建立模型进行训练的时候,可以很容易生成这样的文件,来表示数据:

data/01.jpg,猫
data/05.jpg,狗
data/03.jpg,猫
data/04.jpg,狗
data/06.jpg,狗
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种方法。

如下所示:

images = ...
labels = ...
data = tf.data.Dataset.from_tensor_slices((images, labels))
data = data.batch(batch_size)
iterator = tf.data.Iterator.from_structure(data.output_types,
                                           data.output_shapes)
init_op = iterator.make_initializer(data)
with tf.Session() as sess:
  sess.run(init_op)
  try:
    images, labels = iterator.get_next()
  except tf.errors.OutOfRangeError:
    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读入数据的目的。如下:

def gen():
  with open('train.csv') as f:
    lines = [line.strip().split(',') for line in f.readlines()]
  index = 0
  while True:
    image = cv2.imread(lines[index][0])
    image = cv2.resize(image, (224, 224))
    label = lines[index][1]
    yield (image, label)
    index += 1
    if index == len(lines):
      index = 0


batch_size = 2
data = tf.data.Dataset.from_generator(gen, (tf.float32, tf.int32),
                                      (tf.TensorShape([224, 224, 3]), tf.TensorShape([])))
data = data.batch(batch_size)
iter = data.make_one_shot_iterator()
with tf.Session() as sess:
  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。

示例的代码片段如下:

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

images = tf.constant(image_names)
labels = tf.constant(labels)
images = tf.random_shuffle(images, seed=0)
labels = tf.random_shuffle(labels, seed=0)
data = tf.data.Dataset.from_tensor_slices((images, labels))

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

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

init_op = iterator.make_initializer(data)
with tf.Session() as sess:
  sess.run(init_op)
  try:
    images, filenames, labels = iterator.get_next()
  except tf.errors.OutOfRangeError:
    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的异常出现,需要恰当地进行捕捉并处理。

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

发表评论

您的昵称 *

您的邮箱 *

您的网站