Как установить конкретный gpu в tensorflow?

Я хочу указать gpu для запуска моего процесса. И я поставил его следующим образом:--3--> 

import tensorflow as tf
with tf.device('/gpu:0'):
    a = tf.constant(3.0)
with tf.Session() as sess:
    while True:
        print sess.run(a)

однако он по-прежнему выделять память в два устройства. 

|    0      7479    C   python                         5437MiB 
|    1      7479    C   python                         5437MiB
8
tensorflow tensorflow-gpu

3 ответов

Я считаю, что нужно ставить CUDA_VISIBLE_DEVICES=1. Или какой GPU вы хотите использовать. Если вы сделаете видимым только один GPU, вы будете называть его /gpu:0 независимо от того, что вы установите переменную среды.

дополнительная информация об этой переменной среды:http://www.acceleware.com/blog/cudavisibledevices-masking-gpus

2

есть 3 способа достичь этого:

  1. используя CUDA_VISIBLE_DEVICES переменные среды. установив переменную окруженияCUDA_VISIBLE_DEVICES=1 делает видимым только устройство 1 и, установив CUDA_VISIBLE_DEVICES=0,1 делает устройства 0 и 1 виден.

  2. используя with tf.device('/gpu:2') и создать график. Затем он будет использовать GPU device 2 для запуска.

  3. используя config = tf.ConfigProto(device_count = {'GPU': 1}) а то sess = tf.Session(config=config). Это будет использовать устройство GPU 1.

11

TF выделил бы всю доступную память на каждом видимом GPU, если бы не сказал иначе.

Вы можете установить CUDA_VISIBLE_DEVICES=0,1 если это повторится.

Если вы используете кластер, и вы не хотите отслеживать, какой GPU занят, и вручную введите информацию, которую вы можете вызвать следующий метод перед созданием сеанса. Он будет отфильтровывать графические процессоры, которые уже используются (не имеют много доступной памяти) и установить CUDA_VISIBLE_DEVICES для вы.

функции:

import subprocess as sp
import os

def mask_unused_gpus(leave_unmasked=1):
  ACCEPTABLE_AVAILABLE_MEMORY = 1024
  COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"

  try:
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
    memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
    available_gpus = [i for i, x in enumerate(memory_free_values) if x > ACCEPTABLE_AVAILABLE_MEMORY]

    if len(available_gpus) < leave_unmasked: raise ValueError('Found only %d usable GPUs in the system' % len(available_gpus))
    os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, available_gpus[:leave_unmasked]))
  except Exception as e:
    print('"nvidia-smi" is probably not installed. GPUs are not masked', e)

ограничения: если вы запускаете несколько сценариев одновременно, это все равно вызовет столкновение, потому что память не выделяется сразу при создании сеанса. В случае, если это проблема для вас, вы можете использовать рандомизированную версию, как в моем оригинале исходный код: mask_busy_gpus()

2