web-dev-qa-db-fra.com

Tensorflow définit CUDA_VISIBLE_DEVICES dans jupyter

J'ai deux GPU et souhaite utiliser simultanément deux réseaux différents via ipynb. Cependant, le premier ordinateur portable attribue toujours les deux GPU.

Avec CUDA_VISIBLE_DEVICES, je peux masquer des périphériques pour les fichiers python, mais je ne suis pas sûr de savoir comment le faire dans un cahier.

Est-il possible de cacher différents GPU dans les ordinateurs portables fonctionnant sur le même serveur?

63
Tim

Vous pouvez définir des variables d'environnement dans le bloc-notes à l'aide de os.environ. Effectuez les opérations suivantes avant d’initialiser TensorFlow afin de limiter TensorFlow au premier GPU.

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="0"

Vous pouvez vérifier que les dispositifs appropriés sont visibles par TF.

from tensorflow.python.client import device_lib
print device_lib.list_local_devices()

J'ai tendance à l'utiliser depuis un module utilitaire comme notebook_util

import notebook_util
notebook_util.pick_gpu_lowest_memory()
import tensorflow as tf
123
Yaroslav Bulatov

Vous pouvez le faire plus rapidement sans aucune importation en utilisant simplement magics :

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

Notez que toutes les variables env sont des chaînes, il n'est donc pas nécessaire d'utiliser ". Vous pouvez vérifier que la variable env est configurée en exécutant: %env <name_of_var>. Ou vérifiez-les tous avec %env.

23
Salvador Dali