12

If I call model.cuda() in pytorch where model is a subclass of nn.Module, and say if I have four GPUs, how it will utilize the four GPUs and how do I know which GPUs that are using?

william007
  • 775
  • 1
  • 10
  • 20

1 Answers1

16

model.cuda() by default will send your model to the "current device", which can be set with torch.cuda.set_device(device).

An alternative way to send the model to a specific device is model.to(torch.device('cuda:0')).

This, of course, is subject to the device visibility specified in the environment variable CUDA_VISIBLE_DEVICES.

You can check GPU usage with nvidia-smi. Also, nvtop is very nice for this.

The standard way in PyTorch to train a model in multiple GPUs is to use nn.DataParallel which copies the model to the GPUs and during training splits the batch among them and combines the individual outputs.

noe
  • 26,410
  • 1
  • 46
  • 76
  • If a user has two GPU's, and the 1st GPU is AMD which is not cuda-capable, and the 2nd GPU is NVIDIA, will it still be correct to say model.half().to("cuda:0"), or do I need some logic to check the index (i.e. would the correct input be "cuda:1" in that case and if so how can I detect that?) – pete Apr 05 '22 at 09:22
  • @pete please post your doubt as a new question – noe Apr 05 '22 at 09:27
  • Ok, I asked here https://datascience.stackexchange.com/questions/109682/when-would-i-use-model-tocuda1-as-opposed-to-model-tocuda0 – pete Apr 05 '22 at 10:56