Randomness at the batch level

username: Max_Reynolds
cooked: <p>Hello, I am trying to sample specific parts from a large image (i.e. patches). I would like for the patch indexes (i.e. same parts of images) to be the same within a batch and changed (randomly re-sampled) after every batch.</p>
<p>With num_workers=0, this is straightforward as I can change the state of the dataset after each sample is loaded from the dataloader. However, with multiple workers it seems not as straightforward to alter the workers’ copies of the dataset after each batch.</p>
<p>Is there a good way to make this work with multiple workers? Should I make some kind of custom batch sampler that passes a random patch index sequence to retrieve when sampling? If so, any good resources on customizing the batch sampler?</p>
<p>Simplified example of how this works with num_workers=0</p>
<pre><code class="lang-auto">class exampleDataset(Dataset):
    def __init__(self,dirs,patch_state):
        self.dirs=dirs
        self.patch_state=patch_state
    def __len__(self):
        return len(self.dirs)
    def __getitem__(self,idx):
        image=read_image(dirs[idx])
        patches=get_patches(image,self.patch_state)
    def set_patch_state(self,patch_state):
        self.patch_state=patch_state

initial_patch_state=[0,3,5,33,72]
dataset=exampleDataset(dirs,initial_patch_state)
loader=DataLoader(dataset, batch_size = 8, num_workers=0) 

for i,data in enumerate(loader):
    #&lt;do training&gt;
    #change patches to sample
    new_patch_state=get_5_random_numbers()
</code></pre>
<p>Thanks</p>

username: ptrblck
cooked: <p>If you are using multiple workers, each worker will use a separate seed defined by the main process RNG and the worker id. I guess you could try to use this worker information inside your <code>Dataset.__getitem__</code> to either re-seed the code or change the <code>patch_state</code> accordingly.<br>
To use the states use <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.get_worker_info"><code>torch.utils.data.get_worker_info()</code></a>.</p>

username: Max_Reynolds
cooked: <p>Thank you for the response. I was looking into get_worker_info(), but I think a problem is that there is no concept of the batch in <strong>getitem</strong>. So I’m not sure how you’d be able to only change the state after a batch, rather than after each <strong>getitem</strong> call.</p>

username: ptrblck
cooked: <p>You are right that the default sampler will pass each batch index individually (a <code>BatchSampler</code> could pass all indices of the batch to <code>__getitem__</code>). However, you should also hit the same limitation for <code>num_workers=0</code> so how are you dealing with it now?</p>

username: Max_Reynolds
cooked: <p>With <code>num_workers=0</code>, I can change the patch_state at the end of each patch (as in the example code in the original post above), so that’s how I was handling it. But yes I will proceed with using a BatchSampler; it seems like this is the best approach.</p>
<p><code>__getitem__() </code>will accept a second argument for the patch indexes to return. And within each batch in the batch sampler, the patch indexes will be the same. So in this case there would no longer be a self.patch_state variable for the dataset</p>