How to Get Batch Sample Indices Info in The "__getitem__(self,idx)" Data Loader Instead of Only Getting Single Index Info?

Ardian_Umam: <p>Hi,</p>
<p>I’m currently having a use case of creating custom data loader that can: (i) change <code>batch_size</code> value dynamically during training and (ii) process the data sample with different operation for each different <code>batch_size</code>.</p>
<p>For (i), I successfully accomplish it by creating my custom <a href="https://pytorch.org/docs/stable/data.html" rel="noopener nofollow ugc"><code>batch_sampler</code></a> that returns (a) <code>batch_sample_indices</code> with different number of indices (its <code>len</code> = <code>batch_size</code>). However, I still cannot solve (ii). In order to achieve (ii), I need to know (a), or at least the <code>len</code> of (a) in the <code>__getitem__()</code> method inside the data loader class. Any body can give a direction for this? Or any other approaches to achieve (i) &amp; (ii) is also welcomed.</p>
<p>Thanks in advance!</p>

techkang: <p>You can implement both (i) and (ii) like this:</p>
<pre><code class="lang-auto">import random

from torch.utils.data import DataLoader, Sampler, RandomSampler, Dataset


class MyBatchSampler(Sampler):
    def __init__(self, sampler, batch_size_list, drop_last):
        self.sampler = sampler
        self.batch_size_list = batch_size_list
        self.drop_last = drop_last

    def __iter__(self):
        batch = []
        batch_size = random.choice(self.batch_size_list)
        for idx in self.sampler:
            batch.append((idx, batch_size))
            if len(batch) == batch_size:
                yield batch
                batch = []
                batch_size = random.choice(self.batch_size_list)
        if len(batch) &gt; 0 and not self.drop_last:
            yield batch

    def __len__(self) -&gt; int:
        if self.drop_last:
            return len(self.sampler) // self.batch_size  # type: ignore[arg-type]
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size  # type: ignore[arg-type]


class MyDataSet(Dataset):
    def __getitem__(self, item):
        index, batch_size = item
        return index, batch_size

    def __len__(self):
        return 10


my_dataset = MyDataSet()
my_sampler = MyBatchSampler(RandomSampler(my_dataset), [1, 2, 3], False)
dataloader = DataLoader(my_dataset, batch_sampler=my_sampler)
for data in dataloader:
    print(data)

</code></pre>

Ardian_Umam: <aside class="quote no-group" data-username="techkang" data-post="2" data-topic="135398">
<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/techkang/48/43474_2.png" class="avatar"> techkang:</div>
<blockquote>
<pre><code class="lang-auto">import random

from torch.utils.data import DataLoader, Sampler, RandomSampler, Dataset


class MyBatchSampler(Sampler):
    def __init__(self, sampler, batch_size_list, drop_last):
        self.sampler = sampler
        self.batch_size_list = batch_size_list
        self.drop_last = drop_last

    def __iter__(self):
        batch = []
        batch_size = random.choice(self.batch_size_list)
        for idx in self.sampler:
            batch.append((idx, batch_size))
            if len(batch) == batch_size:
                yield batch
                batch = []
                batch_size = random.choice(self.batch_size_list)
        if len(batch) &gt; 0 and not self.drop_last:
            yield batch

    def __len__(self) -&gt; int:
        if self.drop_last:
            return len(self.sampler) // self.batch_size  # type: ignore[arg-type]
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size  # type: ignore[arg-type]


class MyDataSet(Dataset):
    def __getitem__(self, item):
        index, batch_size = item
        return index, batch_size

    def __len__(self):
        return 10


my_dataset = MyDataSet()
my_sampler = MyBatchSampler(RandomSampler(my_dataset), [1, 2, 3], False)
dataloader = DataLoader(my_dataset, batch_sampler=my_sampler)
for data in dataloader:
    print(data)
</code></pre>
</blockquote>
</aside>
<p>Thanks a lot! Work like a charm <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>