Tensorflow set CUDA_VISIBLE_DEVICES within jupyter

Question

I have two GPUs and would like to run two different networks via ipynb simultaneously, however the first notebook always allocates both GPUs.

Using CUDA_VISIBLE_DEVICES, I can hide devices for python files, however I am unsure of how to do so within a notebook.

Is there anyway to hide different GPUs in to notebooks running on the same server?

Yaroslav Bulatov · Accepted Answer · 2017-02-09T02:10:16.380

181

You can set environment variables in the notebook using os.environ. Do the following before initializing TensorFlow to limit TensorFlow to first GPU.

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

You can double check that you have the correct devices visible to TF

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

I tend to use it from utility module like notebook_util

import notebook_util
notebook_util.pick_gpu_lowest_memory()
import tensorflow as tf

edited Feb 09 '17 at 02:10

answered Jun 18 '16 at 21:43

Yaroslav Bulatov

55,467
20
132
186

4

for me it didn't work os.environ["CUDA_VISIBLE_DEVICES"]="0" ... I changed it to os.environ["CUDA_VISIBLE_DEVICES"]="" and then it worked. I'm using tensorflow 1.12 – Alon Samuel Jun 23 '19 at 15:44
I tried displaying tensorflow local devices as mentioned. My system outputs only device type: CPU. Does this mean that tensorflow is not running GPU. – shaifali Gupta Jul 02 '19 at 13:12
This does not work for me after tensorflow imports. – rjurney Oct 29 '20 at 06:59
That has to be done at the very beginning of your notebook. – muammar Oct 28 '21 at 11:49

score 42 · Answer 2 · answered May 02 '17 at 07:44

42

You can do it faster without any imports just by using magics:

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

Notice that all env variable are strings, so no need to use ". You can verify that env-variable is set up by running: %env <name_of_var>. Or check all of them with %env.

answered May 02 '17 at 07:44

Salvador Dali

199,541
138
677
738

Does this `env` setting affect other python scripts? as it looks like changing the environment variable of os. – JenkinsY May 16 '18 at 11:21
1

@JenkinsY you can easily check it by setting the variable in one script and reading it in another. – Salvador Dali May 23 '18 at 07:35

markemus · Answer 3 · 2021-02-18T17:40:22.407

9

You can also enable multiple GPU cores, like so:

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0,2,3,4"

edited Feb 18 '21 at 17:40

answered Feb 18 '21 at 16:52

markemus

1,294
12
20

Tensorflow set CUDA_VISIBLE_DEVICES within jupyter

3 Answers3

Linked

Related