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

user_3023615 0
Viewed 5457

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?

pytorch deep-learning python-3.x python
2 Answers

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.

user_79963060

To properly implement GPU pre-fetch on PyTorch, you must transform the for-loop into a while-loop.

The DataLoader should be changed into an iterator using the iter function, e.g. iterator = iter(loader).

Use next(iterator) at each step inside the while-loop to get the next mini-batch.

The end of the DataLoader can be detected by catching the StopIteration from the iterator.

Use a flag to end the while-loop when StopIteration is raised.

user_92892750
Related