如果机器上有多于1个的GPU,默认除第一个之外其他的是不工作的,但是可以使用 with ... decice 语句指派特定的CPU或者GPU执行操作。
with tf.Session() as sess: with tf.device('/gpu:1'): a = tf.placeholder(tf.float32) b = tf.placeholder(tf.float32) add = tf.add(a, b)
设备用字符串标识,
cpu:0 机器的CPU gpu:0 机器的第一个GPU,如果有的话 gpu:1 机器的第二个GPU,...
还有通过tf.ConfigProto来构建一个config,在config中指定相关的GPU
tf.ConfigProto参数如下:log_device_placement=True :是否打印设备分配日志 allow_soft_placement=True :如果指定的设备不存在,允许tf自动分配设备 config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True) session = tf.Session(config=config, ......)
控制GPU按需分配
# allow growth config = tf.ConfigProto() config.gpu_options.allow_growth = True # gpu_options=tf.GPUOptions(allow_growth = True) # config=tf.ConfigProto(gpu_options=gpu_options) session = tf.Session(config=config, ...)使用allow_growth option,刚一开始分配少量的GPU容量,然后按需慢慢的增加,由于不会释放
内存,所以会导致碎片
# per_process_gpu_memory_fraction gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.7) config=tf.ConfigProto(gpu_options=gpu_options) session = tf.Session(config=config, ...)设置每个GPU应该拿出多少容量给进程使用,0.4代表 40%
控制使用哪块GPU
~/ CUDA_VISIBLE_DEVICES=0 python your.py # 使用GPU0 ~/ CUDA_VISIBLE_DEVICES=0,1 python your.py # 使用GPU0,1或者在 程序开头
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 使用 GPU 0 os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # 使用 GPU 0,1