I’m working with a dataset where each sample is 120x1024 dimensions. I’m hoping to use a large batch size 4096 (each batch is 4096x120x1024) but am experiencing very slow dataloading even with num_workers=20.
Below is an example of my code. I’m getting 4 iterations/second on my machine. This is a bit too slow to finish training in a reasonable amount of time.
I’ve narrowed the bottleneck down to the torch.stack() call in collate. This seems to be 10x slower than everything else. I believe what is happening is the data is being copied to a contiguous block, hence the slowdown?
Is there any way to speed up dataloading in my case?
class MyDataset(Dataset):
    def __init__(self):
          ...

    def __getitem__(self, index):

         return torch.ones([120, 1024])        


train_dataset = MyDataset()

def collate(batch):
    
    data = torch.stack([b[0] for b in batch])
    
    return data

train_sampler = RandomSampler(train_dataset)

train_dataloader = DataLoader(
    train_dataset, sampler=train_sampler, batch_size=4096,
    num_workers=20, pin_memory=True, collate_fn=collate
)