Tensorflow跨设备调用GPU

最新推荐文章于 2024-07-07 21:47:20 发布

绛洞花主敏明

最新推荐文章于 2024-07-07 21:47:20 发布

阅读量1k

点赞数

分类专栏： tensorflow

本文链接：https://blog.csdn.net/qq_39852676/article/details/108676653

版权

tensorflow 专栏收录该内容

15 篇文章 3 订阅

订阅专栏

注意：在开始讲述之前首先对Tensorflow框架下的模型的跨设备操作进行说明，具体跨设备操作包括以下四种情况：

CPU模型训练 -> GPU模型推理
GPU模型训练 -> GPU模型推理
CPU模型训练 -> CPU模型推理
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.]]