Limit GPU memory consumption

tensorflow

(Claudia) #1

Hello!

I’m running experiments using OpenNMT with tensorflow. I’m running it in a server with 2 GPU’s with the --num_gpus 1 option (because other people need to use the machine too). The process is consuming the memory of both GPU’s, though. Is there any configuration that I can use to limit the amount of memory that the process consume?

Thanks!


(Guillaume Klein) #2

Hello,

You should first restrict the GPUs visible to the process with the CUDA_VISIBLE_DEVICES environment variable, e.g.:

CUDA_VISIBLE_DEVICES=0 onmt-main train [...]

to use GPU 0.

Then you can also set the flag --gpu_allow_growth to only allocate memory as required instead of mapping the entire device memory.