Github下载完整代码
https://github.com/rockingdingo/tensorflow-tutorial/tree/master/mnist
利用tensorflow训练深度神经网络模型需要消耗很长时间,因为并行化计算就为提升运行速度提供了重要思路。Tensorflow提供了多种方法来使程序的并行运行,在使用这些方法时需要考虑的问题有:选取的计算设备是CPU还是GPU,每个CPU多少核的资源并行计算,构建图Graph时消耗资源如何分配等等问题。下面我们以Linux多核CPU的环境为例介绍几种常见方法来提升你的tensorflow程序的运行速度。
在Tensorflow程序中,我们会经常看到”with tf.device("/cpu:0"): “ 这个语句。单独使用这个语句,而不做其他限制,实际上默认tensorflow程序占用所有可以使用的内存资源和CPU核,比如如果你的linux服务器是8核CPU,那么该程序会迅速占用可以使用的任意CPU,使用接近100%,最终结果就是影响整台服务器的其他程序。因此我们会想到需要限制使用的CPU核的个数和资源。
在构建tf.Session() 变量时,可以通过传入tf.ConfigProto() 参数来改变一个tensorflow的session会话所使用的CPU核的个数以及线程数等等。
config = tf.ConfigProto(device_count={"CPU": 4}, # limit to num_cpu_core CPU usage
inter_op_parallelism_threads = 1,
intra_op_parallelism_threads = 1,
log_device_placement=True)
with tf.Session(config = config) as sess:
# To Do
例如,在上面代码中我们通过 “device_count={"CPU": 4}” 参数来构建一个ConfigProto() 类,传入tf.Session()来使每个会话分配相应的资源,这里我们给tensorflow程序共分配了4个CPU core。
我们对tensorflow tutorial中的MNIST的CNN模型convolutional.py进行适当的修改,使得他限制在4个CPU core上运行,得到convolutional_multicore.py: 运行结果,平均每个step计算一个batch的时间为611 ms左右:
图1: 4 CPU core Single Thread
在进行tf.ConfigProto()初始化时,我们也可以通过设置intra_op_parallelism_threads参数和inter_op_parallelism_threads参数,来控制每个操作符op并行计算的线程个数。二者的区别在于:
config = tf.ConfigProto(device_count={"CPU": 4}, # limit to num_cpu_core CPU usage
inter_op_parallelism_threads = 1,
intra_op_parallelism_threads = 4,
log_device_placement=True)
with tf.Session(config = config) as sess:
# To Do
实例比较,下面我们比较线程数为2和4,平均每个batch的运行时间:
总结,在固定CPUcore的资源限制下,通过合理设置线程thread个数可以明显提升tensorflow程序运行速度。
参考 tensorflow/core/protobuf/config.proto 实现
https://github.com/tensorflow/tensorflow/blob/26b4dfa65d360f2793ad75083c797d57f8661b93/tensorflow/core/protobuf/config.proto#L165
有时我们构建的深度网络的结构十分复杂,会出现这种情况:多个CPU core同时运行时,有的核比较空闲,有的核使用率却达到100%的情况。我们需要尽量避免这种运算符计算不均衡的情况。这时,如果我们将Graph拆分为多个部分,将每个部分(如每一层网络结构)指定到不同的CPU 核上运算,优化计算量的分配,可以使运算速度得到提升。
一个很直观的设计就是按照不同的层来划分,把运算量大的Layer分配单独的CPU,把运算量小的Layer合并分配到同一个CPU core上。
下面是我们做的一个测试,还是tensorflow官网上的 convolutional.py 例子改写,将不同层分配到不同的CPU device上,优化了计算资源,使得程序的速度得以提升,例子为convolutional_graph_partitioned.py。
device_id = -1 # Global Variable Counter for device_id used
def next_device(use_cpu = True):
''' See if there is available next device;
Args: use_cpu, global device_id
Return: new device id
'''
global device_id
if (use_cpu):
if ((device_id + 1) < FLAGS.num_cpu_core):
device_id += 1
device = '/cpu:%d' % device_id
else:
if ((device_id + 1) < FLAGS.num_gpu_core):
device_id += 1
device = '/gpu:%d' % device_id
return device
with tf.device(next_device()):
# To Do Insert Your Code
conv = ...
pool = ...