Tensorflow 大规模数据集训练方法

本文转自：Tensorflow】超大规模数据集解决方案：通过线程来预取原文地址：https://blog.csdn.net/mao_xiao_feng/article/details/73991787

现在让我们用Tensorflow实现一个具体的Input pipeline，我们使用CoCo2014作为处理对象，网上应该可以下载到CoCo训练集，train2014这个文件。下载链接：

http://msvocds.blob.core.windows.net/coco2014/train2014.zip

一共13.5G，解压完以后大概会有8万多张图，这个数据集算得上超大规模级别了，那么问题来了，这么多图片我们怎么下手呢？难道和以前一样读到内存？如此笨重的数据集，如果仍然用内存暴力解决，那就太耗费时间空间资源了。能否在训练的同时，读数据，预处理数据呢？现在，让我们用队列+多线程去解决这个问题。

一.Beginning of an input pipeline

在输入pipeline的开始，我们要构造一个队列生成器，tensorflow中的tf.train.string_input_producer函数可以帮助我们解决这个问题：

string_input_producer(string_tensor,num_epochs=None,shuffle=True,seed=None,capacity=32

,shared_name=None,name=None,cancel_op=None)

该函数输入字符串的Tensor或者List，返回一个字符串队列，一共8个输入参数，忽略最后三个不常用的参数，其中

string_tensor：一维的字符串Tensor，注意这个参数不传入Tensor也是可以的，也可以传入字符串的List

num_epochs：控制数量的一个参数，字符串List被放入队列的重复次数

shuffle：很好理解，是否打乱字符串List

seed：shuffle需要的随机数种子，一般可以不指定

capacity：队列的大小

类似的函数还有：

tf.train.range_input_producer

tf.train.slice_input_producer

需要注意的是这里返回的队列是添加了QueueRunner的，也就是我们需要调用线程来操作队列。还有很重要的一点，千万不要以为创建完队列以后，string_tensor的所有值就都入队了，入队也是流程化的，而入队操作通常由分线程来做，任何时刻我们都不关注队列的状态，只关注入队了什么，出队了什么。

二.Batching at the end of an input pipeline

在Input pipeline的末尾则是抽取Batch的流程，这里要用到的函数是tf.train.batch:

batch(tensors,batch_size,num_threads=1,capacity=32,enqueue_many=False

,shapes=None,dynamic_pad=False,allow_smaller_final_batch=False,shared_name=None,name=None)

该函数创建输入的Tensor中的一些batches，同样这个Tensor也可以是一个List，参数解析：

tensors：需要注意的是，为了构成pipeline，保持一致性，这个函数也是以队列形式运行，所以Tensor的输入可以和上面描述的类似。

batch_size：一个batch的数量

num_threads：执行操作的线程数量

capacity：该函数运行的队列的长度

enqueue_many：控制是否可以一次入队多个，一般为false

dynamic_pad：动态填充，填充维度为None的区域，不常用

allow_smaller_final_batch：控制是否允许小于batchsize的batch，如果不允许，则那几个样本会被丢弃

类似函数：

tf.train.batch_join

tf.train.maybe_batch

tf.train.shuffle_batch

同样很重要的一点，batching操作也是在Input pipeline里面，所以他也不是对全部数据来取batch，我们不需要关注当前队列中要多少样本，只需要关注取出了那些样本。

三.示例程序

让我们用一个程序来演示上面的过程。

为了演示需要，从train2014文件夹里面取20张图，然后连着文件夹拷贝到工程路径：

from os import listdir

from os.path import isfile, join

import tensorflow as tf

dataset_path='train2014'

with tf.Session() as sess:

    filenames = [join(dataset_path, f) for f in listdir(dataset_path) if isfile(join(dataset_path, f))]

    print 'number of images:', len(filenames)

    filename_queue = tf.train.string_input_producer(filenames, shuffle=False,num_epochs=1)

    reader = tf.WholeFileReader()

    name, img_bytes = reader.read(filename_queue)

    image = tf.image.decode_jpeg(img_bytes, channels=3)

    dataname = tf.train.batch([name], 2, dynamic_pad=True)

    sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])

    coord = tf.train.Coordinator()

    threads = tf.train.start_queue_runners(coord=coord)

    try:

        while not coord.should_stop():

            print sess.run(dataname)

    except tf.errors.OutOfRangeError:

        print('Done training -- epoch limit reached')

    finally:

        coord.request_stop()

    coord.join(threads)

解决的思路：首先产生一个存储路径下所有文件名的List，然后用它作为输入产生字符串队列，这是pipeline的前端。得到了出队的字符串序列，我们可以使用tensorflow中的filereader，将文件内容读取出来，reader.read函数返回{文件名，文件内容}键值对，文件内容即是我们需要处理的对象（这里为了直观，我们使用了文件名作为输出），这个阶段是pipeline的中端。最后得到了文件名序列，再用batch函数提取一个batch，这个作为pipeline的末端。这样一个程序就完成了。

运行以后结果如下：

number of images: 20

['train2014/COCO_train2014_000000000109.jpg'

 'train2014/COCO_train2014_000000000071.jpg']

['train2014/COCO_train2014_000000000092.jpg'

 'train2014/COCO_train2014_000000000094.jpg']

['train2014/COCO_train2014_000000000064.jpg'

 'train2014/COCO_train2014_000000000025.jpg']

['train2014/COCO_train2014_000000000072.jpg'

 'train2014/COCO_train2014_000000000110.jpg']

['train2014/COCO_train2014_000000000086.jpg'

 'train2014/COCO_train2014_000000000030.jpg']

['train2014/COCO_train2014_000000000113.jpg'

 'train2014/COCO_train2014_000000000009.jpg']

['train2014/COCO_train2014_000000000061.jpg'

 'train2014/COCO_train2014_000000000077.jpg']

['train2014/COCO_train2014_000000000081.jpg'

 'train2014/COCO_train2014_000000000034.jpg']

['train2014/COCO_train2014_000000000078.jpg'

 'train2014/COCO_train2014_000000000036.jpg']

['train2014/COCO_train2014_000000000089.jpg'

 'train2014/COCO_train2014_000000000049.jpg']

Done training -- epoch limit reached

这里我没有shuffle数据集，需要shuffle只要把string_input_producer中的shuffle参数改为True。

四.具体项目

这里没有为大家展示一个具体网络怎么调用以上过程来训练。因为能输出batch的数据其实已经达到我们的意图了。如果有需要一个取数据+训练完整程序的同学，请参考github上图像风格迁移的repo，这个工程使用了以上方法，并有完整的训练过程。源码地址如下：

https://github.com/hzy46/fast-neural-style-tensorflow

---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

下面是原理的讲解部分

一.Tensorflow中的队列机制

队列和线程是Temsorflow中实现异步的重要工具。为什么要异步？用一个形象的例子来解释这个问题。

可以把数据导入的过程看作io操作，在数据规模极大的情况下，io请求需要大量时间执行。同步意味着我们一次处理完io请求，然后再执行程序后面的操作，所以之前在【Tensorflow】怎样为你的网络预加工和打包训练数据？（二）：小数据集的处理方案和【Tensorflow】怎样为你的网络预加工和打包训练数据？（一）中的处理都可以看作同步的，因为都是一次处理完所有的数据，然后在feed给我们的网络。而这里，我们需要扩展一个异步的概念，所谓异步，也就是开辟一个线程来单独处理这个io请求，处理完成，就通知给主程序告诉它“我已经完成了”，在此期间，主程序可以去做其他的事，也就是io请求并不耽误程序的执行。所以异步方式可以显著提高效率，那是不是说异步一定比同步好呢，当然不是，异步只适用于等待时间很长的情况，如果处理小数据集，就不如同步方式了。

Tensorflow是怎样实现异步的？这里需要用到队列这个数据结构，先让我们看看Tensorflow是如何实现队列的。

我们来看一个简单的例子。创建一个“先进先出”队列（FIFOQueue）并填充零。然后，创建一个将元素从队列中取出的图，将该元素加一，并将其放回队列的末尾。缓慢地，队列上的数字增加。用如下的图来表示

和如上图类似，tensorflow中使用RandomShuffleQueue作为输入框架，来准备训练数据。这种机制的运行流程大致如下：

1.使用多线程来准备训练样本，并把他们放入队列

2.一个训练线程执行优化，并吧mini-batch数据提取出队列

TensorFlow Session对象是多线程的，所以多个线程可以轻松地使用相同的Session并并行运行op。但是，实现一个驱动线程的Python程序并不容易。所有线程必须能够一起停止，异常必须被捕获和报告，队列必须在停止时正确关闭。TensorFlow提供了两个类来帮助线程驱动：tf.train.Coordinator和tf.train.QueueRunner。这两个类被设计为一起使用，协调器类帮助多个线程一起停止，并向等待其停止的程序报告异常。QueueRunner类用于创建多个线程，用于协调在同一队列中放入张量。

二.线程协调器

先看一下一些关键方法：

tf.train.Coordinator.should_stop ： returns True if the threads should stop

tf.train.Coordinator.request_stop ：requests that threads should stop

tf.train.Coordinator.join ：waits until the specified threads have stopped

首先创建一个Coordinator对象，然后创建一些使用协调器的线程。线程执行运行循环，当should_stop（）返回True时停止。任何线程都可以决定计算是否停止。它只需要调用request_stop（），其他线程将停止，因为should_stop（）将返回True。

官方也提供了协调器的代码模板：

# Thread body: loop until the coordinator indicates a stop was requested.

# If some condition becomes true, ask the coordinator to stop.

def MyLoop(coord):

  while not coord.should_stop():

    ...do something...

    if ...some condition...:

      coord.request_stop()

# Main thread: create a coordinator.

coord = tf.train.Coordinator()

# Create 10 threads that run 'MyLoop()'

threads = [threading.Thread(target=MyLoop, args=(coord,)) for i in xrange(10)]

# Start the threads and wait for all of them to stop.

for t in threads:

  t.start()

coord.join(threads)

三.队列运行器

QueueRunner类创建了一些重复执行入队操作的线程。这些线程可以使用协调器一起停止，此外，队列运行器运行一个队列关闭的线程，一旦向协调器报告异常，则会自动关闭队列。

创建TensorFlow队列（例如tf.RandomShuffleQueue）用作样本输入的过程如下：

example = ...ops to create one example...

# Create a queue, and an op that enqueues examples one at a time in the queue.

queue = tf.RandomShuffleQueue(...)

enqueue_op = queue.enqueue(example)

# Create a training graph that starts by dequeuing a batch of examples.

inputs = queue.dequeue_many(batch_size)

train_op = ...use 'inputs' to build the training part of the graph...

然后我们要为这个队列创建一个队列运行器，使用多线程来执行入队操作：

# Create a queue runner that will run 4 threads in parallel to enqueue

# examples.

qr = tf.train.QueueRunner(queue, [enqueue_op] * 4)

# Launch the graph.

sess = tf.Session()

# Create a coordinator, launch the queue runner threads.

coord = tf.train.Coordinator()

enqueue_threads = qr.create_threads(sess, coord=coord, start=True)

# Run the training loop, controlling termination with the coordinator.

for step in xrange(1000000):

    if coord.should_stop():

        break

    sess.run(train_op)

# When done, ask the threads to stop.

coord.request_stop()

# And wait for them to actually do it.

coord.join(enqueue_threads)

四.具体的使用方法

在训练当中如何结合以上的工具呢？tf.train.QueueRunner对象要求您在运行任何训练步骤之前调用tf.train.start_queue_runners，否则线程将永久挂起。所以在训练之前要调用tf.train.start_queue_runners启动输入流水线的线程，填充样本队列，以使得样本的出队操作能成功执行。Tensorflow提供了一个参考代码模板：

# Create the graph, etc.

init_op = tf.global_variables_initializer()

# Create a session for running operations in the Graph.

sess = tf.Session()

# Initialize the variables (like the epoch counter).

sess.run(init_op)

# Start input enqueue threads.

coord = tf.train.Coordinator()

threads = tf.train.start_queue_runners(sess=sess, coord=coord)

try:

    while not coord.should_stop():

        # Run training steps or whatever

        sess.run(train_op)

except tf.errors.OutOfRangeError:

    print('Done training -- epoch limit reached')

finally:

    # When done, ask the threads to stop.

    coord.request_stop()

# Wait for threads to finish.

coord.join(threads)

sess.close()

ok，到这里让我们详细解释一下上面的代码可以做什么，首先我们创建计算图，第一阶段获得文件名路径并将其排入文件名队列。第二阶段使用文件名，生成样本，并将它们排列在样本队列中。一旦启动运行这些入队操作的线程，我们的训练循环就可以从样本队列中取出训练样本用以训练。用一张图来表示：

但是在我们实际使用中，要定义一个以上所描述的pipeline是非常复杂的，所以一般就只使用一个文件名队列，而把训练数据的预处理放在主程序当中，因为预处理一个batch的训练数据通常不需要花费太多时间。我们可以使用tf.train.add_queue_runner来添加一个QueueRunner，创建完图后，tf.train.start_queue_runners函数将在图中查询每个QueueRunner，以启动其运行入队操作的线程。此外，我们还要设定epoch的限制，保证在epoch到达时，程序会正确停止，并抛出一个tf.errors.OutOfRangeError异常。

五.一些小的示例程序

首先我们来看第一个先进先出队列例子的实现代码

import tensorflow as tf

q = tf.FIFOQueue(3, "float")

init = q.enqueue_many(vals=[[0., 0., 0.],])

x = q.dequeue()

y = x+1

q_inc = q.enqueue([y])

with tf.Session() as sess:

    init.run()

    q_inc.run()

    q_inc.run()

    q_inc.run()

    q_inc.run()

    print sess.run(x)

    print sess.run(x)

    print sess.run(x)

需要注意的是，enqueue_many方法在实现的时候还是有比较多的坑的，参数vals必须二维，而且要写成[[0,0,0]]的形式而不是，[[0],[0],[0]]，这个要重点记一下。

整个流程的动图上面贴过了，这里再贴一遍吧

然后就是RandomShuffleQueue的例子：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")

sess = tf.Session()

for i in range(10):

    sess.run(q.enqueue('File:'+str(i)))

for i in range(10):

    print(sess.run(q.dequeue()))

因为是RandomShuffleQueue，所以出队顺序是随机的，一直到这里，入队都是用主线程来做的，以上简单版程序没有用多线程。

最后是加上多线程的终极版，让我们来看看：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")

enqueue_op = q.enqueue('File:')

qr = tf.train.QueueRunner(q, enqueue_ops=[enqueue_op] * 1)

sess = tf.Session()

enqueue_threads = qr.create_threads(sess, start=True)

for i in range(100):

　　print(sess.run(q.dequeue()))

好了，运行这个程序，我们会发现输出了100个File，说明在出队的时候，是又有一个线程在做入队操作不停的补充队列的。但是发现一个问题，貌似这个程序自动终止不了？因为我们创建的线程还在运行。

所以要加上协调器，来使得程序得以终止：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")

enqueue_op = q.enqueue('File:')

qr = tf.train.QueueRunner(q, enqueue_ops=[enqueue_op] * 1)

sess = tf.Session()

coord = tf.train.Coordinator()

enqueue_threads = qr.create_threads(sess, start=True,coord=coord)

for i in range(100):

　　print(sess.run(q.dequeue()))

coord.request_stop()

coord.join(enqueue_threads)

这样程序就得以终止了，然后，如果你想再设置一些epoch limit的异常：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")

enqueue_op = q.enqueue('File:')

qr = tf.train.QueueRunner(q, enqueue_ops=[enqueue_op] * 1)

sess = tf.Session()

coord = tf.train.Coordinator()

enqueue_threads = qr.create_threads(sess, start=True,coord=coord)

try:

    for i in range(100):

        print(sess.run(q.dequeue()))

        if i>=50:

            coord.request_stop()

            coord.join(enqueue_threads)

except tf.errors.OutOfRangeError:

    print('Done training -- epoch limit reached')

finally:

    coord.request_stop()

    coord.join(enqueue_threads)