https://www.cnblogs.com/tectal/p/9048184.html
一、TensorFlow 设备分配
1、设备分配规则
If a TensorFlow operation has both CPU and GPU implementations, the GPU devices will be given priority when the operation is assigned to a device.
2、手动指定设备分配
- 如果你不想让系统自动为 operation 分配设备, 而是自己手动指定, 可以用
with tf.device
创建一个设备环境, 这个环境下的 operation 都统一运行在指定的设备上. - 代码示例如下:
-
# op 在 cpu 上运算 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') # op 在 gpu 上运算 with tf.device('/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') # op 在 gpus 上运算 for d in ['/device:GPU:2', '/device: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])
二、TensorFlow GPU 配置
1、指定可以被看见的GPU设备
-
2、限定使用显存的比例import os # 默认情况,TF 会占用所有 GPU 的所有内存, 我们可以指定 # 只有 GPU0 和 GPU1 这两块卡被看到,从而达到限制其使用所有GPU的目的 os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1' # 打印 TF 可用的 GPU print os.environ['CUDA_VISIBLE_DEVICES'] >>> 0, 1
-
3、需要多少拿多少# 在开启对话session前,先创建一个 tf.ConfigProto() 实例对象 # 通过 allow_soft_placement 参数自动将无法放在 GPU 上的操作放回 CPU gpuConfig = tf.ConfigProto(allow_soft_placement=True) # 限制一个进程使用 60% 的显存 gpuConfig.gpu_options.per_process_gpu_memory_fraction = 0.6 # 把你的配置部署到session with tf.Session(config=gpuConfig) as sess: pass 这样,如果你指定的卡的显存是8000M的话,你这个进程只能用4800M。
-
4、GPU 使用总结# 在开启对话session前,先创建一个 tf.ConfigProto() 实例对象 # 通过 allow_soft_placement 参数自动将无法放在 GPU 上的操作放回 CPU gpuConfig = tf.ConfigProto(allow_soft_placement=True) # 运行时需要多少再给多少 gpuConfig.gpu_options.allow_growth = True # 把你的配置部署到session with tf.Session(config=gpuConfig) as sess: pass
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
gpuConfig = tf.ConfigProto(allow_soft_placement=True)
gpuConfig.gpu_options.allow_growth = True
with tf.Session(config=gpuConfig) as sess:
pass