Proper Usage of PyTorch's non_blocking=True for Data Prefetching

Question

I am looking into prefetching data into the GPU from the CPU when the model is being trained on the GPU. Overlapping CPU-to-GPU data transfer with GPU model training appears to require both

1. Transferring data to GPU using data = data.cuda(non_blocking=True)
2. Pin data to CPU memory using train_loader = DataLoader(..., pin_memory=True)

However, I cannot understand how non-blocking transfer is being performed in this official PyTorch example, specifically this code block:

for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        if args.gpu is not None:
            images = images.cuda(args.gpu, non_blocking=True)
        if torch.cuda.is_available():
            target = target.cuda(args.gpu, non_blocking=True)

        # compute output
        output = model(images)
        loss = criterion(output, target)

Won't images.cuda(non_blocking=True) and target.cuda(non_blocking=True) have to be completed before output = model(images) is executed. Since this is a synchronization point, images must be first fully transferred to the CUDA device, so the data transfer steps are effectively no longer non-blocking.

Since output = model(images) is blocking, images.cuda() and target.cuda() in the next i iteration of the for loop will not occur until the model output is computed, meaning no prefetching in the next loop iteration.

If this is correct, what is the correct way to perform data prefetching to the GPU?

Answer 1

I think where you are off is that output = model(images) is a synchronization point. It seems the computation is handled by a different part of a GPU. Quote from official PyTorch docs:

Also, once you pin a tensor or storage, you can use asynchronous GPU copies. Just pass an additional non_blocking=True argument to a to() or a cuda() call. This can be used to overlap data transfers with computation.