If you have installed tensorflow-gpu
and everything runs fine, by default, it should run on your GPU except if:
- You specifically state in your code that the graph should run on CPU with something like
with tf.device('/cpu:0')
,
- You hide your GPU from tensorflow with
os.environ['CUDA_VISIBLE_DEVICES'] = ''
.
You can check that tensorflow finds your GPU using (credit to this answer)
from tensorflow.python.client import device_lib
device_lib.list_local_devices()
Also, if you have an NVIDIA gpu, you can check the usage of your GPU using nvidia-smi
and see that tensorflow is indeed using your GPU.
Now, if you'd like to allocate a specific GPU and/or limit the memory that a graph can allocate on the GPU memory, you might want to have a look here for more details.