Pytorch CPU和GPU并行运行

Question

Is there a way to do something with CPU (compute mean and variance of current mini-batch loss) while GPU is doing back-propagation?

Something like this:

for input, label in dataloader:
  output = model(input)
  losses = some_loss_function(output, label) # size = (batch_size,)
  loss = losses.sum() / batch_size
  # =========== do on CPU ============
  mean = loss.item()
  var = losses.pow(2).sum().item() / batch_size - mean**2
  # ============ BP ================
  loss.backward()
  #gradient update

will the backward() on GPU wait for CPU computation to finish? Is there a way to do backward() and CPU computation in parallel?

Answer 1

It doesn't do both computations at the same time. There is no exposed mechanism to do it parallely. (there is an advanced mechanism using CUDA streams that allows to do this in pytorch, but it is too error-prone for most users)