Tensorflow set CUDA_VISIBLE_DEVICES within jupyter
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"
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
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
.