Asynchronous data loader

In the pytorch, when you build your own dataset and use a pytorch data loader to load dataset. The data loading procedure and networking training procedure and sequentially executed. The execution is shown in Fig. 1.

Fig. 1. Sequential execution

So for one batch training, we need Time T = T1+T2. Sometimes T1 and T2 are both time consuming, e.g. Sometimes you need to pre-process the input data which needs time. Without the increasing of computation resources, can we reduce the Time T? The answer is Yes. Since we are sequentially executed the loading part and training part here, the Time T = T1+T2. If we can parallelly execute the two procedures, the total time T will become roughly T = max(T1, T2). The question is how to do that in pytorch?

The basic idea is using a queue to store the pre-processed data which will be trained soon so the trainer can only get the data from the queue instead of preprocessing first. Here we create a thread separately for the pre-processing task.

Refer to the code below, deliver the dataloader to the CudaDataLoader class, and the device. I use the default queue size of 2 here. If you try to increase the queue size, it will take more GPU memory. If you use only one size queue, it may impact the performance of parallel execution (Sometimes trainer will wait the Loader to load data)

The load_loop function will endless load the data, preprocess the data and put the data into the queue. The load_instance function will load the data to your specified device. The __iter__ and __next__ are also implemented to make the data loader iteratable.

class CudaDataLoader:
    def __init__(self, loader, device, queue_size=2):
        self.device = device
        self.queue_size = queue_size
        self.loader = loader

        self.load_stream = torch.cuda.Stream(device=device)
        self.queue = Queue(maxsize=self.queue_size)

        self.idx =