How to shuffle multi-worker IterableDataset loading

username: MeltingCake
cooked: <p>My data is stored in class-specific files, each file contains all the data for one class.<br>
I’m currently loading it with a custom IterableDataset and a multi-worker DataLoader. As my dataset is quite large, each worker is responsible for a fixed set of <code>n_files/n_workers</code> files as I do not want to every file into memory on every worker.</p>
<p>The problem is as each file is as class-specific, each worker only produces a batch containing the classes it has been assigned (based on worker ID). Each worker has it’s own copy of <code>dataset</code> and <code>collate_fn</code> so it batches within the worker.</p>
<p><a href="https://discuss.pytorch.org/t/how-to-shuffle-an-iterable-dataset/64130" class="inline-onebox">How to shuffle an iterable dataset</a> discusses how to shuffle using <code>torch.utils.data.datapipes.iter.combinatorics.ShuffleIterDataPipe</code> (which isn’t in the docs?). It applies to any iterator but applying it to <code>Dataset</code> then shuffles within each worker and thus batch or to the <code>DataLoader</code> which then only shuffles the order the batches are yielded. The <code>shuffle</code> flag in the <code>DataLoader</code> also just throws me an error about not selecting mode.</p>
<p>Is there some argument in <code>DataLoader</code> that can let me mix and re-collect some buffer of completed batches? Ideally I don’t want to just iterate and collect some <code>n_worker</code> batches and manually shuffle them.</p>
<p>Thanks!</p>
<p>Here’s a simple example of what I’m doing</p>
<pre><code class="lang-auto">class CustomDataset(IterableDataset):
    def __init__(self, files: List[str]):
        self.files = files 
    
    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        
        if worker_info is None:
            files_chunk = self.files
        else:
            n_workers = worker_info.num_workers
            n_files = len(self.files)
            chunk_size = n_files // n_workers
            
            chunk_start = chunk_size * worker_info.id
            files_chunk = self.files[chunk_start: chunk_start + chunk_size]
        
        files_data = [open(f, 'r') for f in files_chunk]
        for line in chain.from_iterable(zip(*files_data)):
            yield line
</code></pre>
<pre><code class="lang-auto">    dataloader = DataLoader(dataset,
                            batch_size=args.batch_size,
                            num_workers=n_cpus,
                            pin_memory=True)
</code></pre>

username: MeltingCake
cooked: <p>I’ve managed to find a solution that works, even if it is a bit ugly.</p>
<pre><code class="lang-auto">  dataloader = DataLoader(dataset,
                          batch_size=1,
                          num_workers=n_cpus,
                          collate_fn=lambda batch: {k:v[0] for k,v in default_collate(batch).items()},
                          prefetch_factor=1)
  shuffled = combinatorics.ShufflerIterDataPipe(dataloader, buffer_size=2 * args.batch_size)
  dataloader = DataLoader(shuffled,
                          batch_size=args.batch_size,
                          num_workers=0,
                          collate_fn=custom_collate_fn)
</code></pre>
<p>The lambda is to reset the extra dimension added to the front when the first dataloader runs with a batch_size=1. <code>ShufflerIterDataPipe</code> then shuffles the single instances before they are batched by the second data loader. In this case, by running it with <code>num_workers=0</code> I can also run GPU operations in there, though I’ve found it to conflict with <code>pin_memory=True</code>.</p>

username: miklos_kralik
cooked: <p>I’ve found an interesting solution to shuffle across batches by subclassing the dataloader:</p>
<p>class BufferShufflingDataLoader(DataLoader):<br>
‘’’ subclasses the pytorch dataloader class to shuffle<br>
the dataset across workers with a buffer size of num workers’‘’</p>
<pre><code>def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.multi_buffer_size = kwargs['num_workers']
    self.multi_buffer = []

def __iter__(self):
    super_iter = super().__iter__()
    for i in range(self.multi_buffer_size):
        self.multi_buffer.append(next(super_iter))

    return self.multi_buffer_iterator(super_iter)

def multi_buffer_iterator(self, super_iter):
    for i in super_iter:
        x = torch.stack(self.multi_buffer)[:, 0]
        _ = self.multi_buffer.pop(0)
        self.multi_buffer.append(next(super_iter))
        yield x
</code></pre>
<p>It is not super efficient in this form bu could be easily sped up by changing the multi buffer list to a tensor of fixed size.</p>