Current CPU device ID?

We can send tensors to different devices. to(device=“cpu:x”)

But how do we get the “cpu:x” device?

The setting is, using MPI to spawn multiple threads for computations (not backprop) and then moving tensors from each thread to a central GPU (for backprop) and back to each CPU.

1 Like

have you considered using nn.dataparallel?

anyway by calling .device you can retrieve the information. Example:

import torch
a=torch.zeros((10,10).cuda()
a.device.index #gpu index
a.device.type # if it is in gpu or cpu

DataParallel does not meet my needs. There is simulator sampling I must do on CPU cores.

a.device.index does not give CPU index. I am looking for similar functionality but for CPU index

1 Like

More broadly, the question can be posed as:
In deep reinforcement learning, if I have multiple CPU cores sampling data from a simulator,
The gradients for these databatches must be passed either serially or synchronously (I’m not sure what the GPU can do) to the GPU, and then the post-gradient-update weights should be sent back to the CPU so that the CPU can use those weights for the new policy to get more samples from the environment.

The problem is I have no idea how to do the above in code. Maybe it needs to be like this:

  1. Every CPU thread holds a copy of the model.
  2. After sampling is done, send the samples and copy of the model to the GPU for backpropagation.
  3. Send model back to CPU.

Step 2 could probably be done by just doing loss.to(device=“cuda:0”) before backpropagation… Hmm… I’ll test this out if there’s no existing examples

I feel like there could be an issue if multiple CPU threads are all sending models to the GPU for backprop and there’s no locking.