Load data in chunks using Dataset

username: naveenkb
cooked: <p>I am wondering if I can modify <code>__get_item__</code> in Dataset to accept multiple indices instead of one index at a time to improve data loading speed from disk using H5 file.</p>
<p>My dataset looks something like this</p>
<pre><code class="lang-auto">class HDFDataset(Dataset):
    def __init__(self, path):
        self.path = path
        
    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        hdf = h5py.File(path, 'r')
        data = hdf['data']
        X = data[idx,:]
        hdf.close()
        return X
</code></pre>
<p><code>dset = HDFDataset(path)</code></p>
<p>I also have a custom batch sampler</p>
<pre><code class="lang-auto">def chunk(indices, chunk_size):
    return torch.split(torch.tensor(indices), chunk_size)


class BatchSampler(Sampler):
    def __init__(self, batch_size, dataset):
        self.batch_size = batch_size
        self.dataset = dataset
    
    def __iter__(self):
        #some function
        return iter(list_of_list_idx)
    
</code></pre>
<p>Sample output from BatchSampler looks as below. Each list represents a batch and each values indicate the indices in a batch</p>
<pre><code class="lang-auto">batch_sampler = BatchSampler(batch_size, dataset)
for x in batch_sampler:
      print(x)
</code></pre>
<blockquote>
<p>[12, 3, 8, 6, 17]<br>
[7, 9, 1, 19, 18]<br>
[13, 4, 2, 5, 14]<br>
[0, 3, 10, 11, 20]</p>
</blockquote>
<p>Dataloader looks like<br>
<code>train_dataloader = DataLoader(dset, num_workers=8, batch_sampler=batch_sampler)</code></p>
<p>This approach works fine for me but data loading takes time as <code>__get_item__</code> loads one index at a time from disk.</p>
<p>Since I already know the indices I need to load in each batch using BatchSampler, is there a way I can load entire batch at once in dataset</p>
<p>for e.g if batch 1 indices is<br>
<code>batch_idx = [12, 3, 8, 6, 17]</code></p>
<p>if <code>__get_item__</code> can accept list of indices than single indice, something like below</p>
<pre><code class="lang-auto">def __getitem__(self, batch_idx):
        hdf = h5py.File(path, 'r')
        data = hdf['data']
        X = data[batch_idx,:]
        hdf.close()
        return X
</code></pre>
<p>Solution 1: Since I already know the indices in each batch I can just load the data to model as tensors, however I won’t be able to utilize num_workers parameter in DataLoader to speed up.</p>
<p>If there is a way to load data in chunks using Dataset &amp; DataLoader, it would solve my issue.<br>
Appreciate any suggestions.</p>

username: ejguan
cooked: <p>In theory, it should work, since fetcher can take a list of tuple indexes.</p><aside class="onebox githubblob">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/780faf52caf170a39414097654d325f6e128414e/torch/utils/data/_utils/fetch.py#L47-L52" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/780faf52caf170a39414097654d325f6e128414e/torch/utils/data/_utils/fetch.py#L47-L52" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/780faf52caf170a39414097654d325f6e128414e/torch/utils/data/_utils/fetch.py#L47-L52</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="47" style="counter-reset: li-counter 46 ;">
          <li>def fetch(self, possibly_batched_index):</li>
          <li>    if self.auto_collation:</li>
          <li>        data = [self.dataset[idx] for idx in possibly_batched_index]</li>
          <li>    else:</li>
          <li>        data = self.dataset[possibly_batched_index]</li>
          <li>    return self.collate_fn(data)</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
For <code>__getitem__</code>, you can take a tuple of indexes to accomplish your request.</p>