How to concatenate different datasets, each with different dimensions

username: Peter_Featherstone
cooked: <p>This is almost exactly the same as this question:</p><aside class="onebox stackexchange">
  <header class="source">

      <a href="https://stackoverflow.com/questions/51837110/pytorch-data-loading-from-multiple-different-sized-datasets" target="_blank" rel="noopener nofollow ugc">stackoverflow.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/10223385/helium4" target="_blank" rel="noopener nofollow ugc">
    <img alt="helium4" src="https://www.gravatar.com/avatar/37562fa318f3a1ce16b46ea0d7e856ee?s=128&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="60" height="60">
  </a>

<h4>
  <a href="https://stackoverflow.com/questions/51837110/pytorch-data-loading-from-multiple-different-sized-datasets" target="_blank" rel="noopener nofollow ugc">PyTorch data loading from multiple different-sized datasets</a>
</h4>

<div class="tags">
  <strong>python, pytorch</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/10223385/helium4" target="_blank" rel="noopener nofollow ugc">
    helium4
  </a>
  on <a href="https://stackoverflow.com/questions/51837110/pytorch-data-loading-from-multiple-different-sized-datasets" target="_blank" rel="noopener nofollow ugc">08:33AM - 14 Aug 18 UTC</a>
</div>

  </article>

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

  <div style="clear: both"></div>
</aside>

<p>I have two datasets A and B. A contains tensors of shape [256,4096] and B contains tensors of shape [32,4096].<br>
Now i can use ConcatDataset to merge A and B, but how do i guarantee that each batch only contains elements from either A or B.<br>
Note, I don’t want to resize elements withing A and B. These are not images.<br>
The answer on stack overflow mentions batch_sampler.<br>
Can somebody elaborate and give a minimal example?</p>

username: naveenkb
cooked: <p>Here is an example of custom batch_sampler for your case</p>
<pre><code class="lang-auto">def chunk(indices, size):
    return torch.split(torch.tensor(indices), size)

class MyBatchSampler(Sampler):
    def __init__(self, a_indices, b_indices, batch_size): 
        self.a_indices = a_indices
        self.b_indices = b_indices
        self.batch_size = batch_size
    
    def __iter__(self):
        random.shuffle(self.a_indices)
        random.shuffle(self.b_indices)
        a_batches  = chunk(self.a_indices, self.batch_size)
        b_batches = chunk(self.b_indices, self.batch_size)
        all_batches = list(a_batches + b_batches)
        all_batches = [batch.tolist() for batch in all_batches]
        random.shuffle(all_batches)
        return iter(all_batches)
   
</code></pre>
<pre><code class="lang-auto">new_dataset = ConcatDataset((dataset_a, dataset_b))
a_len = dataset_a.__len__()
ab_len = a_len + dataset_b.__len__()
a_indices = list(range(a_len))
b_indices = list(range(a_len, b_len))

batch_sampler = MyBatchSampler(a_indices, b_indices, batch_size)

dl = DataLoader(new_dataset,  batch_sampler=batch_sampler)
</code></pre>
<p>to verify if each batch only contains elements from either A or B</p>
<pre><code class="lang-auto">for x in batch_sampler:
    print(x)
</code></pre>

username: Peter_Featherstone
cooked: <p>Thank you. What should <code>def __len__(self)</code> be in <code>MyBatchSampler</code> ?</p>

username: Peter_Featherstone
cooked: <pre><code class="lang-auto">    def __len__(self):
        return (len(self.a_indices) + len(self.b_indices)) // self.batch_size
</code></pre>
<p>??</p>

username: naveenkb
cooked: <p>Yes. You are right. Did it work ?</p>

username: Peter_Featherstone
cooked: <p>Yes it worked thanks</p>