web-dev-qa-db-fra.com

Comment définir un gpu spécifique dans tensorflow?

Je veux spécifier le gpu pour exécuter mon processus. Et je l'ai réglé comme suit:

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)

Cependant, il alloue toujours de la mémoire dans mes deux gpus. 

|    0      7479    C   python                         5437MiB 
|    1      7479    C   python                         5437MiB 
11
luohao

Je crois que vous devez définir CUDA_VISIBLE_DEVICES=1. Ou quel que soit le GPU que vous souhaitez utiliser. Si vous ne rendez visible qu’un seul GPU, vous vous y référerez en tant que /gpu:0 quel que soit le paramétrage de la variable d’environnement.

Plus d'informations sur cette variable d'environnement: https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/

8
Russell

Il y a 3 façons d'y parvenir:

  1. L'utilisation de CUDA_VISIBLE_DEVICES variable d'environnement ..__ en définissant la variable d'environnement CUDA_VISIBLE_DEVICES=1 permet d'afficher uniquement le périphérique 1 et en définissant CUDA_VISIBLE_DEVICES=0,1 les périphériques 0 et 1 sont visibles.

  2. Utiliser with tf.device('/gpu:2') et le créer le graphique. Ensuite, il utilisera le périphérique GPU 2 pour fonctionner.

  3. Utiliser config = tf.ConfigProto(device_count = {'GPU': 1}) puis sess = tf.Session(config=config). Cela utilisera le périphérique GPU 1.

20
Nandeesh

TF allouera toute la mémoire disponible sur chaque GPU visible, sauf indication contraire.

Vous pouvez définir CUDA_VISIBLE_DEVICES=0,1 s'il s'agit d'une chose ponctuelle.

Si vous utilisez un cluster et que vous ne souhaitez pas connaître le GPU occupé ni saisir manuellement les informations, vous pouvez appeler la méthode suivante avant de créer une session. Il filtrera les GPU déjà utilisés (ne disposant pas de beaucoup de mémoire disponible) et définira CUDA_VISIBLE_DEVICES pour vous.

La fonction:

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)

Limitations: si vous démarrez plusieurs scripts à la fois, une collision se produirait, car la mémoire n'est pas allouée immédiatement lorsque vous créez une session. Au cas où cela vous poserait un problème, vous pouvez utiliser une version aléatoire comme dans mon code source original : mask_busy_gpus ()

6
yauheni_selivonchyk

Vous pouvez modifier les paramètres des options du GPU en ajoutant au début de votre script python:

gpu_options = tf.GPUOptions(visible_device_list="0")
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

"0" est ici le nom du GPU que vous voulez utiliser. Vous pouvez avoir la liste du GPU disponible en tapant la commande nvidia-smi dans l'invite du terminal.

0
Tiphaine Champetier