I have a setup with 4 GPUs. When now working with multiple processes in PyTorch, Is there a way to enforce that a process only accesses a given, single gpu, therefore limiting the CUDA driver context to be present only once per process?
Thanks in advance for your help,
Sure, you can do it with the env variable
E.g. to use GPU 0 and 2:
CUDA_VISIBLE_DEVICES=0,2 python pytorch_script.py
and in your case you have to give a different env variable to each process.
Thank you @spanev!
In case anyone is wondering, here is how to set process specific env variables:
import torch.multiprocessing as _mp
mp = _mp.get_context('fork')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
if __name__ == "__main__":
num_processes = 1
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
processes = [Process() for i in range(num_processes)]
[p.start() for p in processes]
print("main: " + os.environ['CUDA_VISIBLE_DEVICES'])
[p.join() for p in processes]
It is important to set it in the run method of the process, as the init method is still called in the main process, therefore setting the env vars of the main process when set there.