2021-22
ssh
to yann.ecs.soton.ac.uk
conda
command (and pip
if neccessary). import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
And, if using Keras follow this with:
import keras.backend.tensorflow_backend as KTF
KTF.set_session(sess)
PyTorch
doesn’t have this problem and will only use what it needs)CUDA_VISIBLE_DEVICES
environment variable; for example, if you want to run on just the first GPU, then CUDA_VISIBLE_DEVICES=0 python myprogram.py
will do the trick. If you have GPU-parallel code, then you can specify multiple GPUs by comma separating IDs: CUDA_VISIBLE_DEVICES=1,2 python myprogram.py
watch nvidia-smi