nn.DataParallel is easy to use when we just have neural network weights.
- What if we have an arbitrary preprocessing (non-differentiable) function in our module? nn.DataParallel does not seem to work well on arbitrary Pytorch tensor functions; at the very least it doesn’t understand how to allocate the tensors dynamically to the right GPU.
For example, I have this normalization code as the preprocessing for my module:
def normalize(self, v): ... return torch.clamp((v.to(torch.cuda.current_device()) - mean.to(torch.cuda.current_device())) / std, -clip_range, clip_range)
I get this error:
RuntimeError: binary_op(): expected both inputs to be on same device, but input a is on cuda:1 and input b is on cuda:0
I get the above error with or without the “to(torch.cuda.current_device())”
-
What if we spawn weights outside of the init function?
E.g. torch.zeros(…) usually gives a similar “Input a is on cuda:1 and input b is on cuda:0” error. I fixed in these cases with “to(torch.cuda.current_device())” but not sure if there’s a silent error affecting the training speed -
How do we actually check that the tensors are being properly distributed across GPUs and we didn’t break the speed somewhere by moving from GPU -> CPU accidentally?
