# Dataloader to load batches into predefined buffer directly

**URL:** <https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511>\
**Category:** Uncategorized\
**Created:** [August 16, 2021, 9:28pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511 "2021-08-16T21:28:15Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![\_iterable](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/32.png) [@\_iterable](https://discuss.pytorch.org/u/_iterable)\
**Post date:** [August 16, 2021, 9:28pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/1 "2021-08-16T21:28:15Z")

</div>

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 \> 0?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [August 17, 2021, 2:50am UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/2 "2021-08-17T02:50:21Z")

</div>

You could take a look at [this post](https://discuss.pytorch.org/t/dataloader-resets-dataset-state/27960/4) showing the usage of `multiprocessing.Array`, which might fit your use case.

---

<div class="post-metadata">

**Author:** ![\_iterable](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/32.png) [@\_iterable](https://discuss.pytorch.org/u/_iterable)\
**Post date:** [August 24, 2021, 10:46pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/3 "2021-08-24T22:46:31Z")

</div>

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.

If my loader looks like this:

```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

```

Is it possible to have a predefined buffer where each worker writes a batch into a position in the buffer?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [August 24, 2021, 11:38pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/4 "2021-08-24T23:38:47Z")

</div>

You could try to create the buffer externally, pass it to the `Dataset`, 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 `Dataset` and pass it as an argument to it.

---

<div class="post-metadata">

**Author:** ![\_iterable](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/32.png) [@\_iterable](https://discuss.pytorch.org/u/_iterable)\
**Post date:** [August 26, 2021, 9:42pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/5 "2021-08-26T21:42:06Z")

</div>

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.

If I have static buffer, can I provide a custom collate function to the datalaoder, which writes the batch into a buffer directly?  
If so, do I need to do special handling for cases where num\_workers \> 0?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [August 26, 2021, 11:55pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/6 "2021-08-26T23:55:54Z")

</div>

> [@\_iterable](#):
>
> 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.

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.

> [@\_iterable](#):
>
> If I have static buffer, can I provide a custom collate function to the datalaoder, which writes the batch into a buffer directly?

I think that’s possible, since my example does the same inside the `Dataset. __getitem__ `.

> [@\_iterable](#):
>
> If so, do I need to do special handling for cases where num\_workers \> 0?

Yes, I also think you would have to reuse my approach of using the shared array.

---

<div class="post-metadata">

**Author:** ![\_iterable](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c67d28/32.png) [@\_iterable](https://discuss.pytorch.org/u/_iterable)\
**Post date:** [August 31, 2021, 4:33pm UTC](https://discuss.pytorch.org/t/dataloader-to-load-batches-into-predefined-buffer-directly/129511/7 "2021-08-31T16:33:02Z")

</div>

I have implemented a version of this:

```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]

```

collator = MyCollator(self.batch\_size, self.in\_height, self.in\_width, 3)  
loader = DataLoader(dataset,  
batch\_size=self.batch\_size,  
drop\_last=True,  
num\_workers=self.num\_workers,  
collate\_fn=collator)

This is however returning the same tensor over multiple batches, what could be the issue here?

Also, each worker gets its own copy of the collator, so there should not be any overwriting of the buffer?
