How does one obtain indicies from a dataloader?

username: Brando_Miranda
cooked: <p>I was looping through the data in a dataloader to create a new dataset. I need to know which data point are which and wanted to get their indicies (e.g. cifar10). How does one do that?</p>
<p>Related:</p>
<p><a href="https://discuss.pytorch.org/t/train-on-a-fraction-of-the-data-set/16743/10">Train on a fraction of the data set</a></p>

username: Brando_Miranda
cooked: <aside class="quote" data-post="1" data-topic="7948">
  <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/f10w/48/1352_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-retrieve-the-sample-indices-of-a-mini-batch/7948">How to retrieve the sample indices of a mini-batch</a> 
  </div>
  <blockquote>
    Hello, 
I would like to know if there is a way to retrieve the indices of the samples in a minibatch: 
for data, target in train_loader:
    # Get the indices of the samples in 'data'
    ...

By sample index I mean the index of that particular sample in the whole original dataset. 
Thank you in advance for your help!
  </blockquote>
</aside>


username: ptrblck
cooked: <p>In the thread you posted is a valid solution:</p><aside class="quote" data-post="2" data-topic="7948">
  <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/ndronen/48/1575_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-retrieve-the-sample-indices-of-a-mini-batch/7948/2">How to retrieve the sample indices of a mini-batch</a> 
  </div>
  <blockquote>
    One way to do this is to implement a subclass of <a href="http://pytorch.org/docs/master/_modules/torch/utils/data/dataset.html#Dataset" rel="nofollow noopener">torch.utils.data.Dataset</a> that returns a triple (data, target, index) from its __getitem__ method. Then your loop would be: 
for data, target, index in train_loader:
    ....
  </blockquote>
</aside>

<p>Here is a small example using CIFAR10 from the other thread:</p><aside class="quote quote-modified" data-post="12" data-topic="7948">
  <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/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-retrieve-the-sample-indices-of-a-mini-batch/7948/12">How to retrieve the sample indices of a mini-batch</a> 
  </div>
  <blockquote>
    You could adapt this code: 
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self):
        self.cifar10 = datasets.CIFAR10(root='YOUR_PATH',
                                        download=False,
                                        train=True,
                                        transform=transforms.ToTensor())
        
    def __getitem__(self, index):
        data, target = self.cifar10[index]
  …
  </blockquote>
</aside>


username: Brando_Miranda
cooked: <p>the thing I can’t generalize is how to do it to any dataloader, they seem to be specific to the corresponding data sets, or am I missing something?</p>

username: ptrblck
cooked: <p>The index is specific to a <code>Dataset</code> and you can return it in the <code>__getitem__</code> function.<br>
The <code>DataLoader</code> just calls the <code>__getitem__</code> function from its <code>Dataset</code> and iterates it using the specified batch size.</p>
<p>I don’t think there is an easy way to modify a <code>DataLoader</code> to return the index. At least, I don’t have an idea, sorry.</p>

username: Brando_Miranda
cooked: <pre><code>def __getitem__(self, index):
    data, target = self.cifar10[index]
    return data, target, index</code></pre>

username: arturml
cooked: <p>Hi <a class="mention" href="/u/brando_miranda">@Brando_Miranda</a>, did you solve your problem already? When I saw your last question what came into my mind was to extract the data from cifar10 as a numpy array, like this</p>
<pre><code class="lang-python">data = cifar10.train_data
</code></pre>
<p>Checkout the source code <a href="https://github.com/pytorch/vision/blob/master/torchvision/datasets/cifar.py" rel="nofollow noopener">here</a>. Now you can work with this array, use your rule to get the indices you want and pass them do <code>SubsetRandomSampler</code>.</p>

username: Brando_Miranda
cooked: <p>Thanks! I wish that would have occurred to me to look into…</p>

username: cassidylaidlaw
cooked: <p>See here for a general solution: <a href="https://discuss.pytorch.org/t/how-to-retrieve-the-sample-indices-of-a-mini-batch/7948/19" class="inline-onebox">How to retrieve the sample indices of a mini-batch</a></p>

username: simonxxxx
cooked: <p>Here my implementation.</p>
<pre><code class="lang-auto">from torch.utils.data import DataLoader, Dataset

class IndexDataset(Dataset):
    def __init__(self, dataset, subset=None):
        self.dataset = dataset
        self.subset = subset
        
    def __getitem__(self, index):
        if self.subset is None:
            data = self.dataset[index]
            real_index = index
        else:
            real_index = self.subset[index]
            data = self.dataset[real_index]
            
        if isinstance(data, dict):
            data["real_index"] = real_index
            return data
        elif isinstance(data, list):
            return [real_index] + data
        else:
            raise NotImplementedError(f"Data type {type(data)} not supported")

    def __len__(self):
        if self.subset is not None:
            return len(self.subset)
        else:
            return len(self.dataset)

</code></pre>