Asynchronous processing of GPU and CPU in pyg

[RX28666]

Hello,

I am currently running a model with two iterations of forward/backward passes, each with a different batch size, on a single GPU. The process is currently sequential: (1) sample one batch with batch size 1, (2) sample another batch with batch size 2, (3) execute the model with the first sampled batch, and (4) execute the model with the second sampled batch. Each step waits for the previous one to be completed. The code looks like:

kwargs_1 = {'batch_size': 1024, 'num_workers': 16, 'persistent_workers': True}
train_loader1 = NeighborLoader(data, input_nodes=None,
                              num_neighbors=[-1, -1], shuffle=True, **kwargs_1)

kwargs_2 = {'batch_size': 2048, 'num_workers': 16, 'persistent_workers': True}
train_loader2 = NeighborLoader(data, input_nodes=None,
                              num_neighbors=[-1, -1], shuffle=True, **kwargs_2)

 dataloader_iterator = iter(train_loader2)
 for i, batch1 in enumerate(train_loader1):
       try:
           batch2 = next(dataloader_iterator)
       except StopIteration:
           dataloader_iterator = iter(train_loader2)
           batch2 = next(dataloader_iterator)
           
       optimizer.zero_grad()
       out1 = model(batch1)
       out2 = model(batch2)
        .......

However, since sampling on the CPU takes longer than executing the forward and backward passes on the GPU, I am exploring the possibility of running them iteratively. For example, during step 2 (sampling with batch size 2) on the CPU, I want to concurrently run the model with batch size 1 on the GPU (step 3) in an asynchronous way. If this is possible, how can I achieve it in pyg? Thanks!

[rusty1s]

I think this should already do the trick:

for i, batch1 in enumerate(train_loader1):
       optimizer.zero_grad()
       out1 = model(batch1)
       try:
           batch2 = next(dataloader_iterator)
       except StopIteration:
           dataloader_iterator = iter(train_loader2)
           batch2 = next(dataloader_iterator)
           
       out2 = model(batch2)
        .......