天天看点

用GPU进行TensorFlow计算加速

小编说:将深度学习应用到实际问题中,一个非常大的问题在于训练深度学习模型需要的计算量太大。为了加速训练过程,本文将介绍如何如何在TensorFlow中使用单个GPU进行计算加速,也将介绍生成TensorFlow会话(tf.Session)时的一些常用参数。通过这些参数可以使调试更加方便而且程序的可扩展性更好。本文选自《TensorFlow:实战Google深度学习框架(第2版)》。

TensorFlow程序可以通过tf.device函数来指定运行每一个操作的设备,这个设备可以是本地的CPU或者GPU,也可以是某一台远程的服务器。但在本文中只关心本地的设备。TensorFlow会给每一个可用的设备一个名称,tf.device函数可以通过设备的名称来指定执行运算的设备。比如CPU在TensorFlow中的名称为/cpu:0。在默认情况下,即使机器有多个CPU,TensorFlow也不会区分它们,所有的CPU都使用/cpu:0作为名称。而一台机器上不同GPU的名称是不同的,第n个GPU在TensorFlow中的名称为/gpu:n。比如第一个GPU的名称为/gpu:0,第二个GPU名称为/gpu:1,以此类推。

TensorFlow提供了一个快捷的方式来查看运行每一个运算的设备。在生成会话时,可以通过设置log_device_placement参数来打印运行每一个运算的设备。以下程序展示了如何使用log_device_placement这个参数。

在以上代码中,TensorFlow程序生成会话时加入了参数log_device_placement=True,所以程序会将运行每一个操作的设备输出到屏幕。于是除了可以看到最后的计算结果,还可以看到类似“add:

/job:localhost/replica:0/task:0/cpu:0”这样的输出。这些输出显示了执行每一个运算的设备。比如加法操作add是通过CPU来运行的,因为它的设备名称中包含了/cpu:0。

在配置好GPU环境的TensorFlow中,如果操作没有明确地指定运行设备,那么TensorFlow会优先选择GPU。比如将以上代码在亚马逊(Amazon Web Services, AWS)的 g2.8xlarge实例上运行时,会得到类似以下的运行结果。

从以上输出可以看到在配置好GPU环境的TensorFlow中,TensorFlow会自动优先将运算放置在GPU上。不过,尽管g2.8xlarge实例有4个GPU,在默认情况下,TensorFlow只会将运算优先放到/gpu:0上。于是可以看见在以上程序中,所有的运算都被放在了/gpu:0上。如果需要将某些运算放到不同的GPU或者CPU上,就需要通过tf.device来手工指定。以下程序给出了一个通过tf.device手工指定运行设备的样例。

在以上代码中可以看到生成常量a和b的操作被加载到了CPU上,而加法操作被放到了第二个GPU“/gpu:1”上。在TensorFlow中,不是所有的操作都可以被放在GPU上,如果强行将无法放在GPU上的操作指定到GPU上,那么程序将会报错。以下代码给出了一个报错的样例。

不同版本的TensorFlow对GPU的支持不一样,如果程序中全部使用强制指定设备的方式会降低程序的可移植性。在TensorFlow的kernel中定义了哪些操作可以跑在GPU上。比如可以在variable_ops.cc程序中找到以下定义。

在这段定义中可以看到GPU只在部分数据类型上支持tf.Variable操作。如果在TensorFlow代码库中搜索调用这段代码的宏TF_CALL_GPU_NUMBER_TYPES,可以发现在GPU上,tf.Variable操作只支持实数型(float16、float32和double)的参数。而在报错的样例代码中给定的参数是整数型的,所以不支持在GPU上运行。为避免这个问题,TensorFlow在生成会话时可以指定allow_soft_placement参数。当allow_soft_placement参数设置为True时,如果运算无法由GPU执行,那么TensorFlow会自动将它放到CPU上执行。以下代码给出了一个使用allow_soft_placement参数的样例。

虽然GPU可以加速TensorFlow的计算,但一般来说不会把所有的操作全部放在GPU上。一个比较好的实践是将计算密集型的运算放在GPU上,而把其他操作放到CPU上。GPU是机器中相对独立的资源,将计算放入或者转出GPU都需要额外的时间。而且GPU需要将计算时用到的数据从内存复制到GPU设备上,这也需要额外的时间。TensorFlow可以自动完成这些操作而不需要用户特别处理,但为了提高程序运行的速度,用户也需要尽量将相关的运算放在同一个设备上。

TensorFlow默认会占用设备上的所有GPU以及每个GPU的所有显存。如果在一个TensorFlow程序中只需要使用部分GPU,可以通过设置CUDA_VISIBLE_DEVICES环境变量来控制。以下样例介绍了如何在运行时设置这个环境变量。

TensorFlow也支持在程序中设置环境变量,以下代码展示了如何在程序中设置这些环境变量。

虽然TensorFlow默认会一次性占用一个GPU的所有显存,但是TensorFlow也支持动态分配GPU的显存,使得一块GPU上可以同时运行多个任务。下面给出了TensorFlow动态分配显存的方法。

<a target="_blank" href="https://item.jd.com/12287533.html"></a>

用GPU进行TensorFlow计算加速

继续阅读