PyTorch等深度学习框架在使用多个GPU时如何处理内存？

Question

I have recently run into a situation where I am running out of memory on a single Nvidia V100. I have limited experience using multiple GPUs to train networks so I'm a little unsure on how the data parallelization process works. Lets say I'm using a model and batch size that requires something like 20-25GB of memory. Is there any way to take advantage of the full 32GB of memory I have between two 16GB V100s? Would PyTorch's DataParallel functionality achieve this? I suppose there is also the possibility of breaking the model up and using model parallelism as well. Please excuse my lack of knowledge on this subject. Thanks in advance for any help or clarification!

Answer 1

You should keep model parallelism as your last resource and only if your model doesn't fit in the memory of a single GPU (with 16GB/GPU you have plenty of room for a gigantic model).

If you have two GPUs, I would use data parallelism. In data parallelism you have a copy of your model on each GPU and each copy is fed with a batch. The gradients are then gathered and used to update the copies.

Pytorch makes it really easy to achieve data parallelism, as you just need to wrap you model instance in nn.DataParallel :

model = torch.nn.DataParallel(model, device_ids=[0, 1])
output = model(input_var)