Dataloader to load batches into predefined buffer directly

_iterable: <p>Can the pytorch dataloader be configured to write data into a predefined buffer? Instead of dynamically allocating memory for each batch, each batch could be written into a buffer and read from there. Is this possible currently? Can this be done even when num_workers &gt; 0?</p>

ptrblck: <p>You could take a look at <a href="https://discuss.pytorch.org/t/dataloader-resets-dataset-state/27960/4">this post</a> showing the usage of <code>multiprocessing.Array</code>, which might fit your use case.</p>

_iterable: <p>Thank you for your suggestion! I tried this but this does not directly solve my use case. What I would want is essentially the Dataloader to not dynamically create a tensor for each batch, but write each batch into a predefined buffer.</p>
<p>If my loader looks like this:</p>
<pre><code class="lang-auto">loader = DataLoader(
    dataset,
    num_workers=7,
    shuffle=False
)
loader_iter = iter(loader)
buffer  # size of this is 2*num_workers

next(loader_iter) # this should write the new batch into the buffer directly
</code></pre>
<p>Is it possible to have a predefined buffer where each worker writes a batch into a position in the buffer?</p>

ptrblck: <p>You could try to create the buffer externally, pass it to the <code>Dataset</code>, and write the data to it from each worker (making sure that the workers don’t overlap). I guess the difference between my approach and your suggestion would be to initialize the shared array outside of the <code>Dataset</code> and pass it as an argument to it.</p>

_iterable: <p>My dataset is too huge to hold the entire data in the buffer, so I would like to hold only intermediate batches returned by the dataloader in the buffer.</p>
<p>If I have static buffer, can I provide a custom collate function to the datalaoder, which writes the batch into a buffer directly?<br>
If so, do I need to do special handling for cases where num_workers &gt; 0?</p>

ptrblck: <aside class="quote no-group" data-username="_iterable" data-post="5" data-topic="129511">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/48.png" class="avatar"> _iterable:</div>
<blockquote>
<p>My dataset is too huge to hold the entire data in the buffer, so I would like to hold only intermediate batches returned by the dataloader in the buffer.</p>
</blockquote>
</aside>
<p>Would you like to create a caching mechanism for e.g. 10% of the data and reuse it in the next epoch? If you are filling (and clearing) the cache with new samples in each iteration, I would expect to see mostly misses, since your would basically implement a running window which doesn’t contain the next sample.</p>
<aside class="quote no-group" data-username="_iterable" data-post="5" data-topic="129511">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/48.png" class="avatar"> _iterable:</div>
<blockquote>
<p>If I have static buffer, can I provide a custom collate function to the datalaoder, which writes the batch into a buffer directly?</p>
</blockquote>
</aside>
<p>I think that’s possible, since my example does the same inside the <code>Dataset.__getitem__</code>.</p>
<aside class="quote no-group" data-username="_iterable" data-post="5" data-topic="129511">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/48.png" class="avatar"> _iterable:</div>
<blockquote>
<p>If so, do I need to do special handling for cases where num_workers &gt; 0?</p>
</blockquote>
</aside>
<p>Yes, I also think you would have to reuse my approach of using the shared array.</p>

_iterable: <p>I have implemented a version of this:</p>
<pre><code class="lang-auto">MyCollator:
def __init__(self, batch_size, height, width, channels):
        self.buffer = torch.zeros(8, self.batch_size, self.channels, self.height, self.width, dtype=torch.bfloat16)
        self.curr_idx = -1

def __call__(self, batch):
        elem = batch[0]
        elem_type = type(elem)

        if isinstance(elem, torch.Tensor):
            self.curr_idx = (self.curr_idx + 1) % 8
            torch.stack(batch, 0, out=self.buffer[self.curr_idx])
            return self.buffer[self.curr_idx]
</code></pre>
<p>collator = MyCollator(self.batch_size, self.in_height, self.in_width, 3)<br>
loader = DataLoader(dataset,<br>
batch_size=self.batch_size,<br>
drop_last=True,<br>
num_workers=self.num_workers,<br>
collate_fn=collator)</p>
<p>This is however returning the same tensor over multiple batches, what could be the issue here?</p>
<p>Also, each worker gets its own copy of the collator, so there should not be any overwriting of the buffer?</p>