TensorFlow使用并行计算

物理CPU、逻辑CPU、CPU核数、多线程与并行运算

在学习如何使用TensorFlow并行计算之前,我们必须要明白一些相关概念。这将有助于我们更好的学习和理解TensorFlow的并行运算机制。

多个物理CPU

这句话的意思是一台计算机有多个物理CPU

物理cpu数:主板上实际插入的cpu数量,可以数不重复的 physical id 有几个(physical id)

cpu核数

单块CPU上面能处理数据的芯片组的数量,如双核、四核等 (cpu cores)

逻辑cpu数

一般情况下,逻辑cpu=物理CPU个数×每颗核数,如果不相等的话,则表示服务器的CPU支持超线程技术(HT:简单来说,它可使处理器中的1 颗内核如2 颗内核那样在操作系统中发挥作用。这样一来,操作系统可使用的执行资源扩大了一倍,大幅提高了系统的整体性能,此时逻辑cpu=物理CPU个数×每颗核数x2)

多线程

我们知道线程是操作系统最小的调度单位,单个线程同时只能在单个cpu线程(CPU线程相当于逻辑CPU数)中执行,对于多核的CPU系统中,开启多线程有助于加速程序的运行。

并行运算

并行计算(英语:parallel computing)一般是指许多指令得以同时进行的计算模式。在同时进行的前提下,可以将计算的过程分解成小部分,之后以并发方式来加以解决。多线程也算做并行运算。一般用以加速程序的运算。

TensorFlow中使用并行运算

在TensorFlow中使用并行运算有很多方式,比如增加线程、使用多块GPU、使用多块CPU等,因为我使用的服务器上只有两块CPU,所以下面我们以使用CPU为例,GPU也是一样的。

  • 手动指定网络的某一部分运行的CPU
with tf.device('/cpu:0'):
    要运行的代码块...

with tf.device('/cpu:1'):
    要运行的代码块...
...
  • 使用TensorFlow提供的配置API设置线程池
config = tf.ConfigProto(device_count={"CPU": num_cpu_core}, 
            inter_op_parallelism_threads = x, 
            intra_op_parallelism_threads = x,
            log_device_placement=True)

with tf.Session(config=config) as sess:
    代码块...

上面参数中的两种配置是通过调整线程池来优化 CPU 性能:

  • intra_op_parallelism_threads: 使用多线程来并行化计算的结点会将不同的计算单元分配到该池中的线程上
  • inter_op_parallelism_threads: 所有待计算结点都由此线程池来调度
  • 值得注意的
    TensorFlow官方文档写到,对于以上两种配置,如果都没有设置或设置为 0,则会默认使用逻辑 CPU 核心的数目。对于许多系统,包括 4 核的 CPU,以及包含
    70 多个逻辑核心的多 CPU 系统,测试都显示默认设置已经非常高效。

关于TensorFlow使用并行运算,解释起来比较抽象,我们还是用一个例子更好的来说明一下吧!

使用并行运算加速MNIST数据集的训练

首先描述一下本例子使用到的神经网络:我们搭建了一个CNN神经网络,两层卷积层,两层全连接层。
网络的相关参数设置如下:

  • batch_size=128
  • epoch_size=1000

使用TensorFlow提供的配置API设置线程池

配置设置如下:

config = tf.ConfigProto(device_count={"CPU": 4}, 
                inter_op_parallelism_threads = 2, 
                intra_op_parallelism_threads = 4,
                log_device_placement=True)

运行结果:

配置运行过程使用的核心数和所开线程数的方式加速:
Extracting ./data/mnist/train-images-idx3-ubyte.gz
Extracting ./data/mnist/train-labels-idx1-ubyte.gz
Extracting ./data/mnist/t10k-images-idx3-ubyte.gz
Extracting ./data/mnist/t10k-labels-idx1-ubyte.gz
0.0761
0.8259
0.893
0.9184
0.9262
0.9391
0.945
0.9497
0.9544
0.955
0.9575
0.959
0.9616
0.9652
0.9669
0.9697
0.9684
0.9705
0.9717
0.9709
batch_time: 0.20636706495285034
peak memory: 6262.14 MiB, increment: 6042.41 MiB
CPU times: user 16min 14s, sys: 2min 7s, total: 18min 22s
Wall time: 5min 7s

可见,每训练一个batch耗时batch_time=0.20636706495285034s; 总耗时5min7s

切分网络,使其运行在不同的CPU上,再配合线程的配置

配置设置如下:

with tf.device('/cpu:0'):
        数据加载
        卷积层
        
    with tf.device('/cpu:0'):
        全连接层
        梯度计算
    
    config = tf.ConfigProto(device_count={"CPU": 4}, 
                inter_op_parallelism_threads = 2, 
                intra_op_parallelism_threads = 4,
                log_device_placement=True)

运行结果:

Extracting ./data/mnist/train-images-idx3-ubyte.gz
Extracting ./data/mnist/train-labels-idx1-ubyte.gz
Extracting ./data/mnist/t10k-images-idx3-ubyte.gz
Extracting ./data/mnist/t10k-labels-idx1-ubyte.gz
0.1135
0.7231
0.8665
0.9071
0.9238
0.9322
0.9389
0.9442
0.9501
0.9519
0.9537
0.957
0.958
0.962
0.965
0.9658
0.964
0.9683
0.9693
0.9698
batch_time: 0.19959420514106752
peak memory: 6239.59 MiB, increment: 6020.55 MiB
CPU times: user 15min 48s, sys: 2min, total: 17min 49s
Wall time: 4min 58s

可见,每训练一个batch耗时batch_time=0.19959420514106752; 总耗时4min 58s,比第一种方式快。

再一次调整线程配置

配置设置如下:

config = tf.ConfigProto(device_count={"CPU": 20}, 
                inter_op_parallelism_threads = 2, 
                intra_op_parallelism_threads = 40,
                log_device_placement=True)

运行结果:

配置运行过程使用的核心数和所开线程数的方式加速:
Extracting ./data/mnist/train-images-idx3-ubyte.gz
Extracting ./data/mnist/train-labels-idx1-ubyte.gz
Extracting ./data/mnist/t10k-images-idx3-ubyte.gz
Extracting ./data/mnist/t10k-labels-idx1-ubyte.gz
0.124
0.8032
0.8919
0.9134
0.9293
0.9391
0.9461
0.95
0.9521
0.9556
0.9581
0.9609
0.9646
0.964
0.9652
0.9656
0.9675
0.9694
0.9707
0.9702
batch_time: 0.053559075355529784
peak memory: 6271.99 MiB, increment: 6052.40 MiB
CPU times: user 27min 6s, sys: 6min 47s, total: 33min 53s
Wall time: 1min 22s

可见,每训练一个batch耗时batch_time=0.053559075355529784; 总耗时1min 22s,比第第一个配置方式快了近4倍。

综述

通过上面的实验,我们可以看出通过配置网络使用不同的硬件,增加对运算的线程和使用到的CPU核心数等,可以大大提高神经网络的执行效率,节约训练时间。