如果机器上有多于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
GitHub 加速计划 / te / tensorflow
184.54 K
74.12 K
下载
一个面向所有人的开源机器学习框架
最近提交(Master分支:1 个月前 )
a49e66f2 PiperOrigin-RevId: 663726708 1 个月前
91dac11a This test overrides disabled_backends, dropping the default value in the process. PiperOrigin-RevId: 663711155 1 个月前
Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐