TensorFlow中的多线程

正文共1699个字，2张图，预计阅读时间5分钟。

TensorFlow提供两个类帮助实现多线程，一个是tf.train.Coordinator，另一个是tf.train.QueueRunner。Coordinator主要用来实现多个线程同时停止，QueueRunner用来创建一系列线程。

Coordinator

根据官方文档，Coordinator主要有三个方法：

1、tf.train.Coordinator.should_stop: returns True if the threads should stop.

2、tf.train.Coordinator.request_stop: requests that threads should stop.

3、tf.train.Coordinator.join: waits until the specified threads have stopped.

接下来我们实验Coordinator，下面的代码主要实现每个线程独立计数，当某个线程达到指定值的时候，所有线程终止：

#encoding=utf-8

import threading

import numpy as np

import tensorflow as tf

#创建一个函数实现多线程，参数为Coordinater和线程号

def func(coord, t_id):    

count = 0    

while not coord.should_stop(): #不应该停止时计数        

print('thread ID:',t_id, 'count =', count)        

count += 1        

if(count == 5): #计到5时请求终止            

coord.request_stop()

coord = tf.train.Coordinator()

threads = [threading.Thread(target=func, args=(coord, i)) for i in range(4)]

#开始所有线程

for t in threads:    

t.start()

coord.join(threads) #等待所有线程结束

运行结果如下，当0号线程打印出4时，其他线程不再计数，程序终止。

QueueRunner

QueueRunner的作用是创建一些重复进行enqueue操作的线程，它们通过coordinator同时结束。

#encoding=utf-8

import numpy as np

import tensorflow as tf batch_size = 2

#随机产生一个2*2的张量

example = tf.random_normal([2,2])

#创建一个RandomShuffleQueue，参数意义参见API

q = tf.RandomShuffleQueue(    capacity=1000,    min_after_dequeue=0,    dtypes=tf.float32,    shapes=[2,2])

#enqueue op，每次push一个张量

enq_op = q.enqueue(example)

#dequeue op, 每次取出batch_size个张量

xs = q.dequeue_many(batch_size)

#创建QueueRunner，包含4个enqueue op线程

qr = tf.train.QueueRunner(q, [enq_op]*4) coord = tf.train.Coordinator() sess = tf.Session()

#启动QueueRuner，开始线程

enq_threads = qr.create_threads(sess, coord=coord, start=True)for i in range(10):    

if coord.should_stop():        

break    print('step:', i, sess.run(xs))

#打印结果

coord.request_stop() coord.join(enq_threads)

总结

这两个类是实现TensorFlow pipeline的基础，能够高效地并行处理数据。个人认为在数据较大时，应该避免使用feed_dict。因为，feed_dict是利用python读取数据，python读取数据的时候，tensorflow无法计算，而且会将数据再次拷贝一份。

原文链接：https://www.jianshu.com/p/c29965e6c40c

查阅更为简洁方便的分类文章以及最新的课程、产品信息，请移步至全新呈现的“LeadAI学院官网”：

www.leadai.org

请关注人工智能LeadAI公众号，查看更多专业文章

大家都在看

LSTM模型在问答系统中的应用

基于TensorFlow的神经网络解决用户流失概览问题

最全常见算法工程师面试题目整理（一）

最全常见算法工程师面试题目整理（二）

TensorFlow从1到2 | 第三章 深度学习革命的开端：卷积神经网络

装饰器 | Python高级编程

今天不如来复习下Python基础