1. 指定 GPU 运算
如果安装的是 GPU 版本, 在运行的过程中 TensorFlow 能够自动检测. 如果检测到 GPU,TensorFlow 会尽可能的利用找到的第一个 GPU 来执行操作.
如果机器上有超过一个可用的 GPU, 除了第一个之外的其他的 GPU 默认是不参与计算的. 为了让 TensorFlow 使用这些 GPU, 必须将 OP 明确指派给他们执行. with......device 语句能够用来指派特定的 CPU 或者 GPU 执行操作:
- import tensorflow as tf
- import numpy as np
- with tf.Session() as sess:
- with tf.device('/cpu:0'):
- a = tf.placeholder(tf.int32)
- b = tf.placeholder(tf.int32)
- add = tf.add(a, b)
- sum = sess.run(add, feed_dict={a: 3, b: 4})
- print(sum)
设备的字符串标识, 当前支持的设备包括以下的几种:
CPU:0 机器的第一个 CPU.
gpu:0 机器的第一个 gpu, 如果有的话
gpu:1 机器的第二个 gpu, 依次类推
类似的还有 tf.ConfigProto 来构建一个 config, 在 config 中指定相关的 GPU, 并且在 session 中传入参数 config="自己创建的 config" 来指定 gpu 操作
其中, tf.ConfigProto 函数的参数如下:
log_device_placement=True: 是否打印设备分配日志
allow_soft_placement=True: 如果指定的设备不存在, 允许 TF 自动分配设备
- import tensorflow as tf
- import numpy as np
- config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)
- with tf.Session(config=config) as sess:
- a = tf.placeholder(tf.int32)
- b = tf.placeholder(tf.int32)
- add = tf.add(a, b)
- sum = sess.run(add, feed_dict={a: 3, b: 4})
- print(sum)
2. 设置 GPU 使用资源
上文的 tf.ConfigProto 函数生成的 config 之后, 还可以设置其属性来分配 GPU 的运算资源, 如下代码就是按需分配
- import tensorflow as tf
- import numpy as np
- config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)
- config.gpu_options.allow_growth = True
- with tf.Session(config=config) as sess:
- a = tf.placeholder(tf.int32)
- b = tf.placeholder(tf.int32)
- add = tf.add(a, b)
- sum = sess.run(add, feed_dict={a: 3, b: 4})
- print(sum)
使用 allow_growth option, 刚开始会分配少量的 GPU 容量, 然后按需要慢慢的增加, 有与不会释放内存, 随意会导致内存碎片.
同样, 上述的代码也可以在 config 创建时指定,
- import tensorflow as tf
- import numpy as np
- gpu_options = tf.GPUOptions(allow_growth=True)
- config = tf.ConfigProto(gpu_options=gpu_options)
- with tf.Session(config=config) as sess:
- a = tf.placeholder(tf.int32)
- b = tf.placeholder(tf.int32)
- add = tf.add(a, b)
- sum = sess.run(add, feed_dict={a: 3, b: 4})
- print(sum)
我们还可以给 gpu 分配固定大小的计算资源.
gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.5)
上述代码的含义是分配给 tensorflow 的 GPU 显存大小为: GPU 的实际显存 * 0.5
来源: https://www.cnblogs.com/baby-lily/p/10924444.html