博客

Tensorflow跨设备调用GPU

  • Post author:
  • Post category:其他



注意

：在开始讲述之前首先对

Tensorflow

框架下的模型的跨设备操作进行说明，具体跨设备操作包括以下四种情况：


  1. CPU

    模型训练 ->

    GPU

    模型推理

  2. GPU

    模型训练 ->

    GPU

    模型推理

  3. CPU

    模型训练 ->

    CPU

    模型推理

  4. GPU

    模型训练 ->

    CPU

    模型推理

以上四种情况在

Pytorch

框架下存在差异，需要在不同的设备上（例如

GPU

训练->

CPU

推理的代码要在模型加载中加入

模型.to(device)

）进行模型的加载，但在

Tensorflow

框架下不存在对模型的改变，只需要对默认的设备进行修改即可。不需要单独声明模型的加载设备。

  • 安装的

    tensorlow-gpu

    ，且设备中存在

    GPU

    可调用，设备默认执行在

    GPU

    中。
  • 安装的

    tensorlow-cpu

    ，设备默认执行在

    CPU

    中。



1、准备

在进行

GPU

的调用和测试之前首先要确保使用的机器上存在

GPU

设备，即在可使用

GPU

的基础下进行环境的搭建。具体工作包括如下：


  • CUDA

    、

    Cudnn

    、显卡驱动等基本环境的安装。

  • Tensorflow-GPU

    的安装。

搭建完环境之后对环境进行测试，即是否能成功检测到

GPU

，以及能够成功进行调用。



2、查看设备是否支持GPU

首先进入安装

Tensorflow-gpu

的环境当中，输入命令python，并进入

Python

环境。

在这里插入图片描述

正常情况下输入如下测试代码，即可打印出机器上存在的可用设备，测试代码如下： 

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

例如设备上存在一块

GPU

和

CPU

，成功执行后的代码如下： 

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 13177083330855175469
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 10968950375
locality {
  bus_id: 1
}
incarnation: 6161624703599064583
physical_device_desc: "device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:00:08.0, compute capability: 6.1"
]

从上述执行打印的代码可以看出，执行设备的名称，通过在

Tensorflow

中指定设备的名称可以实现对代码的设备管理。

在一套标准的系统上通常有多个计算设备.

TensorFlow

支持

CPU

和

GPU

这两种设备. 我们用指定字符串

strings

来标识这些设备. 比如:

  • “

    /cpu:0

    ”: 机器中的

    CPU
  • “

    /gpu:0

    ” 或 ”

    /device:GPU:0

    ”: 机器中的

    GPU

    , 如果你有一个的话.
  • “

    /gpu:1

    ” 或 ”

    /device:GPU:1

    ”: 机器中的第二个

    GPU

    , 以此类推…

如果一个

TensorFlow

的

operation

中兼有

CPU

和

GPU

的实现, 当这个算子被指派设备时,

GPU

有优先权. 比如

matmul

中

CPU

和

GPU kernel

函数都存在. 那么在

cpu:0

和

gpu:0

中,

matmul operation

会被指派给

gpu:0

.



3、记录设备指派情况

为了获取你的

operations

和

Tensor

被指派到哪个设备上运行, 用

log_device_placement

新建一个

session

, 并设置为

True

. 

# 新建一个 graph.
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# 新建session with log_device_placement并设置为True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 运行这个 op.
print sess.run(c)

你应该能看见以下输出: 

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
id: 0000:05:00.0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
MatMul: /job:localhost/replica:0/task:0/gpu:0
[[ 22.  28.]
 [ 49.  64.]]



4、使用OS模块声明可见设备

有一台服务器，服务器上有多块儿

GPU

可以供使用，但此时只希望使用第2块和第4块

GPU

，但是我们希望代码能看到的仍然是有两块

GPU

，分别编号为

0,1

，这个时候我们可以使用环境变量

CUDA_VISIBLE_DEVICES

来解决这个问题。

比如： 

CUDA_VISIBLE_DEVICES=1  只有编号为1的GPU对程序是可见的，在代码中gpu[0]指的就是这块儿GPU
CUDA_VISIBLE_DEVICES=0,2,3  只有编号为0,2,3的GPU对程序是可见的，在代码中gpu[0]指的是第0块儿，gpu[1]指的是第2块儿，gpu[2]指的是第3块儿
CUDA_VISIBLE_DEVICES=2,0,3  只有编号为0,2,3的GPU对程序是可见的，但是在代码中gpu[0]指的是第2块儿，gpu[1]指的是第0块儿，gpu[2]指的是第3块儿



5、allow_soft_placement参数自动切换设备

例如，一台服务器上存在

GPU

，代码默认使用

GPU

进行调用，但模型的代码程序中存在只可在

CPU

上运行的程序，那么此时代码会报错，此时

allow_soft_placement

设置为

True

，则在GPU设备不可用的情况下，默认将此代码切换为CPU进行运行，此时的工程中存在一部分代码在GPU中进行运行，一部分代码在CPU中运行。使用方法如下： 

with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:

tf.ConfigProto用来对session进行参数配置，如上所示。涉及到的参数如下： 

#tf.ConfigProto()的参数
log_device_placement=True : 是否打印设备分配日志
allow_soft_placement=True ： 如果你指定的设备不存在，允许TF自动分配设备

除外，还可以控制GPU的使用率等，如下： 

gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.7)#设置每个GPU使用率0.7代表70%
config=tf.ConfigProto(gpu_options=gpu_options)
session = tf.Session(config=config, ...)



6、手动指派设备

如果你不想使用系统来为 operation 指派设备, 而是手工指派设备, 你可以用 with tf.device 创建一个设备环境, 这个环境下的 operation 都统一运行在环境指定的设备上. 

# 新建一个graph.
with tf.device('/cpu:0'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
  b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# 新建session with log_device_placement并设置为True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 运行这个op.
print sess.run(c)

你会发现现在 a 和 b 操作都被指派给了 cpu:0. 

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
id: 0000:05:00.0
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
MatMul: /job:localhost/replica:0/task:0/gpu:0
[[ 22.  28.]
 [ 49.  64.]]


（1）在多GPU系统里使用单一GPU，

如果你的系统里有多个 GPU, 那么 ID 最小的 GPU 会默认使用. 如果你想用别的 GPU, 可以用下面的方法显式的声明你的偏好: 

# 新建一个 graph.
with tf.device('/gpu:2'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
  b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
  c = tf.matmul(a, b)
# 新建 session with log_device_placement 并设置为 True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 运行这个 op.
print sess.run(c)

如果你指定的设备不存在, 你会收到

InvalidArgumentError

错误提示: 

InvalidArgumentError: Invalid argument: Cannot assign a device to node 'b':
Could not satisfy explicit device specification '/gpu:2'
   [[Node: b = Const[dtype=DT_FLOAT, value=Tensor<type: float shape: [3,2]
   values: 1 2 3...>, _device="/gpu:2"]()]]

为了避免出现你指定的设备不存在这种情况, 你可以在创建的 session 里把参数 allow_soft_placement 设置为 True, 这样 tensorFlow 会自动选择一个存在并且支持的设备来运行 operation. 

# 新建一个 graph.
with tf.device('/gpu:2'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
  b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
  c = tf.matmul(a, b)
# 新建 session with log_device_placement 并设置为 True.
sess = tf.Session(config=tf.ConfigProto(
      allow_soft_placement=True, log_device_placement=True))
# 运行这个 op.
print sess.run(c)


（2）使用多个 GPU

如果你想让 TensorFlow 在多个 GPU 上运行, 你可以建立 multi-tower 结构, 在这个结构 里每个 tower 分别被指配给不同的 GPU 运行. 比如: 

# 新建一个 graph.
c = []
for d in ['/gpu:2', '/gpu:3']:
  with tf.device(d):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
    c.append(tf.matmul(a, b))
with tf.device('/cpu:0'):
  sum = tf.add_n(c)
# 新建session with log_device_placement并设置为True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 运行这个op.
print sess.run(sum)

你会看到如下输出: 

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K20m, pci bus
id: 0000:02:00.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: Tesla K20m, pci bus
id: 0000:03:00.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: Tesla K20m, pci bus
id: 0000:83:00.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: Tesla K20m, pci bus
id: 0000:84:00.0
Const_3: /job:localhost/replica:0/task:0/gpu:3
Const_2: /job:localhost/replica:0/task:0/gpu:3
MatMul_1: /job:localhost/replica:0/task:0/gpu:3
Const_1: /job:localhost/replica:0/task:0/gpu:2
Const: /job:localhost/replica:0/task:0/gpu:2
MatMul: /job:localhost/replica:0/task:0/gpu:2
AddN: /job:localhost/replica:0/task:0/cpu:0
[[  44.   56.]
 [  98.  128.]]


版权声明：本文为qq_39852676原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。