Correct use of collate_fn

username: bornexmachina
cooked: <p>My data has a condition for each sample I process. During the training I would like to batch through my data, such that each batch only contains data from one condition.</p>
<p>Something like</p>
<ol start="0">
<li>for batch size of n</li>
<li>pick a condition randomly</li>
<li>pick n elements from the condition</li>
</ol>
<p>My custom Data Set has the following return in</p>
<pre><code class="lang-auto">__getitem__(self, index):
    ...
    return data, y_true, segment, basename, index, operating_condition
</code></pre>
<p>Then I define</p>
<pre><code class="lang-auto">def collate_fn(batch):
    *_, conditions = zip(*batch)
    unique_conditions = list(set(conditions))
    condition_batches = {c: [] for c in unique_conditions}

    for b in batch:
        condition_batches[b[5]].append(b)

    batches = [condition_batches[c] for c in unique_conditions]
    return batches
</code></pre>
<p>which returns a<br>
<code>[[list of samples of condition A], [list of samples of condition B], .., ]</code></p>
<p>finally I use torch’s DataLoader as</p>
<pre><code class="lang-auto">self.train_loader = torch.utils.data.DataLoader(
            self.train_dataset, collate_fn=collate_fn,
            batch_size=batch_size, shuffle=shuffle, batch_sampler=batch_sampler,
        )
</code></pre>
<p>yet when I print out my batch which I get by</p>
<pre><code class="lang-auto">batch = next(iter(train_loader))
print(batch)
</code></pre>
<p>I see again a mix of conditions in my batch.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a25562bc51f7d37095e3db7f5697c06ae9ea64da.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/a25562bc51f7d37095e3db7f5697c06ae9ea64da" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/2/a25562bc51f7d37095e3db7f5697c06ae9ea64da_2_690x433.jpeg" alt="image" data-base62-sha1="na4hRQsslv21iQ0GrVvxNELNlTc" width="690" height="433" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/2/a25562bc51f7d37095e3db7f5697c06ae9ea64da_2_690x433.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/a/2/a25562bc51f7d37095e3db7f5697c06ae9ea64da_2_1035x649.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/2/a25562bc51f7d37095e3db7f5697c06ae9ea64da_2_1380x866.jpeg 2x" data-dominant-color="ECE7ED"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1866×1173 145 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: bornexmachina
cooked: <p>Well, I have investigated / worked more on this.</p>
<p>Here is a small working example with my mock data structure.</p>
<pre><code class="lang-auto">import numpy as np

b = (np.array([ 1.17, 1.65, 1.61, 7.61]), 0.0, np.array([1.]), 'file_vel_16_loc_A.wav', 206050, 'vel_16_loc_A')
c = (np.array([ 2.17, 2.65, 2.61, 2.61]), 0.0, np.array([1.]), 'file_vel_8_loc_D.wav', 206050, 'vel_8_loc_D')

batch = [b, c, b, c, b]

def collate_fn(batch):
    *_, conditions = zip(*batch)
    unique_conditions = list(set(conditions))
    condition_batches = {c: [] for c in unique_conditions}

    for b in batch:
        condition_batches[b[5]].append(b)

    batches = [list(zip(*condition_batches[c])) for c in unique_conditions]
    for i in range(len(batches)):
        for j in [0, 1, 2, 4]:
            batches[i][j] = np.stack(batches[i][j], axis=0)
    return batches


out = collate_fn(batch)
</code></pre>
<p>and my out looks like intended. It is a list of length 2, where each element is grouped by the condition.<br>
So far so good.</p>
<pre><code class="lang-auto">[
# here comes condition one called vel_8_loc_D
[array([[2.17, 2.65, 2.61, 2.61],
         [2.17, 2.65, 2.61, 2.61]]),
  array([0., 0.]),
  array([[1.], [1.]]),
  ('file_vel_8_loc_D.wav', 'file_vel_8_loc_D.wav'),
  array([206050, 206050]),
  ('vel_8_loc_D', 'vel_8_loc_D')],

# and now the condition two vel_16_loc_A
 [array([[1.17, 1.65, 1.61, 7.61],
         [1.17, 1.65, 1.61, 7.61],
         [1.17, 1.65, 1.61, 7.61]]),
  array([0., 0., 0.]),
  array([[1.],
         [1.],
         [1.]]),
  ('file_vel_16_loc_A.wav', 'file_vel_16_loc_A.wav', 'file_vel_16_loc_A.wav'),
  array([206050, 206050, 206050]),
  ('vel_16_loc_A', 'vel_16_loc_A', 'vel_16_loc_A')]
]
</code></pre>
<p>Now, when I iterate through my batches each of the batches is a list. Length of the list is equal to the unique conditions found in the batch.</p>
<p>So basically what happens is:</p>
<ol>
<li>A random batch is selected with batch size of e.g. 256</li>
<li>This 256 samples are grouped based on their conditions. E.g. 50 samples in group A, 100 samples in group B, 106 samples in group C.</li>
<li>Now I have three lists where I could pick some samples from</li>
</ol>
<p>Unfortunately, it is not exactly what I would like.</p>
<p>If I transfer it to an stochastic urn model I would do the following, assuming I need a batch of 256 elements</p>
<ol>
<li>Pick first ball from the urn and note its color. E.g. red</li>
<li>Pick further balls from the urn.<br>
2a. If the ball is any different color, place it back in the urn<br>
2b. If the ball is red add it to the sample.</li>
<li>Repeat 2a and 2b until you have 256 red balls</li>
</ol>
<p>From the documentation of the <code>collate_fn</code> I deduce, that this approach is NOT feasible, as <code>collate_fn</code><br>
comes into action after I have picked 256 random balls, and it just sorts them by color.</p>
<p>Can anybody suggest, how to model my problem in Torch?</p>
<p>In theory, I could just have arbitrary large batch sizes and then hope to get enough samples in the particular element of the list. But it is obviously hacky and unreliable.</p>
<p>I have implemented  a dictionary of DataLoaders, each only for a specific condition. BUT for my model I would need always a batch from each condition where I know, that it is this particular condition and a random query, which might come from any of the potential conditions DataLoaders.</p>
<p>So if following the approach of condition specific DataLoaders, how do I ensure, that I actually processed through all of the samples with my random query?</p>
<p>In my toy example:</p>
<p>Imagine I have 3 conditions, each with 100 samples. So in total 300 samples. I wish to make 300 queries so each sample is definitely covered. And each of the samples will be compared with a randomly drawn sample from condition 1 only, condition 2 only, condition 3 only. So I will cycle 3 times through the condition specific  DataLoader but only once through the general DataLoader.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="bornexmachina" data-post="2" data-topic="181520">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bornexmachina/48/58727_2.png" class="avatar"> bornexmachina:</div>
<blockquote>
<p>In theory, I could just have arbitrary large batch sizes and then hope to get enough samples in the particular element of the list. But it is obviously hacky and unreliable.</p>
</blockquote>
</aside>
<p>Also, it would add a massive overhead of loading a huge number of samples where the majority might be wasted.</p>
<aside class="quote no-group" data-username="bornexmachina" data-post="2" data-topic="181520">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bornexmachina/48/58727_2.png" class="avatar"> bornexmachina:</div>
<blockquote>
<p>Can anybody suggest, how to model my problem in Torch?</p>
</blockquote>
</aside>
<p>Your use case sounds as if you might want to implement a custom <code>sampler</code>, which is responsible to create the indices passed to the <code>Dataset</code>, not a custom <code>collate_fn</code>, which is responsible to create a batch using the already loaded samples.</p>