Why not move images onto gpu?

[I-Doctor]

examples/imagenet/main.py

Line 284 in 792d336

if args.gpu is not None:

I'm trying to training vgg on imagenet with one node DataParallel and no multiprocessing。But I find 'images.device' before computation is 'cpu', and 'target.device=cuda:0'. I'm not sure why these four lines of codes move 'images' to gpu only when I choose only one gpu(args.gpu is not None) and move 'target' to gpu even with argument device=None(args.gpu=None).

I would appreciate it if someone could help me understand it.

[subramen]

Thanks @I-Doctor, this is a weird snippet that has multiple occurrences in this script. The second condition (if torch.cuda.is_available()) can be merged into the first (if args.gpu is not None). Would you like to raise a PR for this?

[soumith]

what @suraj813 suggested is incorrect.

I'm not sure why these four lines of codes move 'images' to gpu only when I choose only one gpu(args.gpu is not None) and move 'target' to gpu even with argument device=None(args.gpu=None).

the reason is this.
When you do single GPU training, you dont use nn.DataParallel, you simply send images directly to your model. So images is moved onto GPU 0 in a straight-forward manner.

When you do multi-GPU training with nn.DataParallel, you keep images on CPU and then send it to nn.DataParallel, which will slice images into as many Tensors as there are GPUs (for example if images had 64 images and there are 4 GPUs, images is sliced into four separate Tensors of 16 images each). These individual sliced Tensors are then sent to their respective GPUs. In the 4-GPU example, if we had sent images onto GPU 0 and then called nn.DataParallel there would be a wasteful move of images first to GPU0, and then each of the 3 slices of images will then each be moved from GPU0 to GPU1, GPU0 to GPU2 and GPU0 to GPU3 respectively.

[subramen]

Thanks @soumith for clarifying what that if block is doing. IIUC the two conditions are to cater to both nn.DataParallel and DistributedDataParallel

What I had missed was that inputs don't need to be on GPU, nn.DataParallel manages moving it to GPUs under the hood

x = torch.rand(4,1)  # on CPU
model = nn.DataParallel(model)  # replicated on all GPUs
y = model(x)  # gathered on GPU0

So when args.gpu == None (i.e. multi-GPU training with nn.DataParallel), the first if block keeps images on CPU. The second if block moves targets to GPU0 in case of DataParallel, or to the respective rank if using DistributedDataParallel.

Closing the issue, @I-Doctor please reopen if you have follow-up queries!