I am working with a dataset where samples take uneven amounts of time to load. Most samples load on the order of 20ms while some samples take much longer (i.e. 10 seconds). I have a custom Dataset and I’m using a DataLoader to parallelize the loading process.
I’ve noticed that even when I set shuffle=True, the data loader will block waiting for certain workers. In an optimal situation, you would expect all available workers to be processing an “expensive” sample at any given moment but what I observe is that the workers are underutilized. A worker that processed a fast sample may be blocked on a worker that is processing a slow sample.
I can implement my own solution using multiprocessing.Pool.imap_unordered but I wanted to see if there is a canonical PyTorch way to solve this problem.