How to retrieve the sample indices of a mini-batch

f10w: <p>Hello,</p>
<p>I would like to know if there is a way to retrieve the indices of the samples in a minibatch:</p>
<pre><code>for data, target in train_loader:
    # Get the indices of the samples in 'data'
    ...
</code></pre>
<p>By sample index I mean the index of that particular sample in the whole original dataset.<br>
Thank you in advance for your help!</p>

ndronen: <p>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"><code>torch.utils.data.Dataset</code></a> that returns a triple <code>(data, target, index)</code> from its <code>__getitem__</code> method. Then your loop would be:</p>
<pre><code class="lang-auto">for data, target, index in train_loader:
    ....
</code></pre>

f10w: <p>Thanks, Nicholas. Did you mean <code>DataLoader</code> instead of <code>Dataset</code>?</p>
<p>My <code>train_loader</code> is a DataLoader, so <code>data</code> is a mini-batch and not a single sample. I guess I should modify  <code>DataLoader</code>? Sorry I’m a newbie <img src="//discuss.pytorch.org/images/emoji/apple/frowning.png?v=5" title=":frowning:" class="emoji" alt=":frowning:"></p>

smth: <p>he meant Dataset. The returned 3rd argument from the DataLoader will be the indices of the mini-batch.</p>

f10w: <p>Many thanks to <a class="mention" href="/u/ndronen">@ndronen</a> and <a class="mention" href="/u/smth">@smth</a>!!</p>

vithursant: <p>Did you end up solving your problem? I would also like retrieve the indices of the samples in a mini-batch.</p>

ptrblck: <p>The suggested solution should work. Is it not working in your case?</p>

vithursant: <p>Yes, I got it working afterwards. It was a bug on my end. But yes, the suggested solution works. Thanks!</p>

CalebEverett: <p>Here’s an example class for the STL10 data set. I’d be happy to add the indexes to the datasets and submit a pull request if that would be useful.</p>
<pre><code>class STL10DataSet(Dataset):
"""
STL10 data set. Implemented to be able to get the index of the images
in a mini-batch so that predictions can be associated with their original
image data.

https://github.com/mttk/STL10/blob/master/stl10_input.py
"""

def __init__(self, root_dir, split='train', transform=None):
    """
    Args:
        root_dir (string): path where stl10_binary.tar.gz has been extracted.
        transform (callable, optional): Optional transform to be applied on an
            image.
    """
    self.root_dir = root_dir
    self.split = split
    self.transform = transform
    
    with open(os.path.join(self.root_dir, '{}_y.bin'.format(self.split)), 'rb') as f:
        self.labels = np.fromfile(f, dtype=np.uint8) - 1
        
    with open(os.path.join(self.root_dir, 'class_names.txt')) as f:
        self.classes = f.read().splitlines()
        
def __len__(self):
    return len(self.labels)

def __getitem__(self, idx):
    with open(os.path.join(self.root_dir, '{}_X.bin'.format(self.split)), 'rb') as f:
        c, h, w = 3, 96, 96
        size = c*h*w
        f.seek(idx*size)
        # read bytes from file
        image = np.fromfile(f, dtype=np.uint8, count=size)
        # reshape to C x H x W
        image = np.reshape(image, (c, h, w))
        # transpose to H x W x C for ToPILImage transform
        image = np.transpose(image, (2, 1, 0))
        
        if self.transform:
            image = self.transform(image)
            
        return image, self.labels[idx], idx</code></pre>

Brando_Miranda: <p>can you provide the import statement for Dataset?</p>

Brando_Miranda: <p>sorry Im confused, how does you code change if I wanted to use cifar10?</p>

ptrblck: <p>You could adapt this code:</p>
<pre><code class="lang-auto">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]
        
        # Your transformations here (or set it in CIFAR10)
        
        return data, target, index

    def __len__(self):
        return len(self.cifar10)

dataset = MyDataset()
loader = DataLoader(dataset,
                    batch_size=1,
                    shuffle=True,
                    num_workers=1)

for batch_idx, (data, target, idx) in enumerate(loader):
    print('Batch idx {}, dataset index {}'.format(
        batch_idx, idx))
</code></pre>

Brando_Miranda: <p>I think its important to note that:</p>
<pre><code>data, target = self.cifar10[index]
</code></pre>
<p>returns floatTensor and not PIL stuff, so one needs to cast to it or numpy for the transforms to work I believe. Thanks for the help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>That’s right! You can remove the <code>transform=transforms.ToTensor()</code> line when creating the dataset, if you would like to apply PIL transformation in <code>__getitem__</code>.<br>
Or just add it in the constructor with <code>transforms.Compose</code>. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=5" title=":wink:" class="emoji" alt=":wink:"></p>

Brando_Miranda: <p>what if I told u I had a bunch of neural nets creating the data set <img src="https://discuss.pytorch.org/images/emoji/apple/rofl.png?v=5" title=":rofl:" class="emoji" alt=":rofl:"><img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"></p>
<p>wish I was joking</p>

simo23: <p>Hi,</p>
<p>if for anyone else the already posted solutions are not enough:</p>
<ol>
<li>
<p>In torch.utils.data.Dataloader.py in the function “put_indices” add this line at the end of the function:<br>
<code>return indices</code></p>
</li>
<li>
<p>In the same file, in the function right below “put_indices” called “_process_next_batch” modify the line:<br>
<code>self._put_indices()</code><br>
to be:<br>
<code>indices = self._put_indices() # indices contains the indices in the batch.</code></p>
</li>
<li>
<p>To have the indices available while iterating the dataloader modify the line in the same file, same function “_process_next_batch” from:<br>
<code>return batch</code><br>
To:<br>
<code>return (batch, indices)</code></p>
</li>
<li>
<p>You can now know the indices in the training script with:</p>
</li>
</ol>
<pre><code class="lang-auto">     for data in dataloaders['train']:
         (input, target), indices = data
</code></pre>
<p>The code in the “Dataloader” should now be like this:</p>
<pre><code class="lang-auto">    def _put_indices(self):
        assert self.batches_outstanding &lt; 2 * self.num_workers
        indices = next(self.sample_iter, None)
        if indices is None:
            return
        self.index_queues[self.worker_queue_idx].put((self.send_idx, indices))
        self.worker_queue_idx = (self.worker_queue_idx + 1) % self.num_workers
        self.batches_outstanding += 1
        self.send_idx += 1
        return indices # Added this line

    def _process_next_batch(self, batch):
        self.rcvd_idx += 1
        # self._put_indices() # Old line
        indices = self._put_indices() # Modified line
        if isinstance(batch, ExceptionWrapper):
            raise batch.exc_type(batch.exc_msg)
        #return batch # Old line
        return (batch, indices) # Modified line
</code></pre>
<ol start="5">
<li>To know the correspondences between indices and images you can simply add a line in your training script. I have my datasets defined in a standard way as:</li>
</ol>
<p><code>image_datasets = {x: datasets.ImageFolder(os.path.join(args.data_dir, x),data_transforms[x]) for x in ['train', 'val']}</code></p>
<p>So to know the index of all the images in the “train” dataset use:</p>
<p><code>image_datasets['train'].samples</code></p>
<p>The output is a list containing the paths of the images and their classes ids. The order is <strong>alphabetical</strong> and it is obtained by first sorting the folders of the dataset (classes) and then the files inside each folder (images).</p>
<p><strong>Note</strong>: to know where the files of a certain module that you are <strong>currently</strong> using  are you can type:</p>
<pre><code class="lang-auto">python
import torch.utils.data
torch.utils.data.dataloader.__file__
</code></pre>

azunino: <p>Hi. <a class="mention" href="/u/simo23">@simo23</a> thank you for your proposed solution, but it seems not properly work. If the imgs are prepared in a shuffle way the idx do not correspond to the correct original images. I have tried to not shuffle the data images in my train loader and see the idx of the images of the first batch. They are in order but start from 80, they should start from 0 <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=6" title=":frowning:" class="emoji" alt=":frowning:"> (e.g if my batch size is 10, the obtained idx are 80,81,… 89). This problem seems also related to the number of workers used in the train loader. If I am using num of workers = 4 the idx start from 80 if I am using num of workers = 1 they start from 20. Is there anyone who has the same problem?</p>

cassidylaidlaw: <p>If anyone is still looking for an easy solution, this is what I did. This function takes a dataset <em>class</em> (not an instance) and returns a class that returns a data, target, index tuple from <code>__getitem__</code>:</p>
<pre><code>def dataset_with_indices(cls):
"""
Modifies the given Dataset class to return a tuple data, target, index
instead of just data, target.
"""

    def __getitem__(self, index):
        data, target = cls.__getitem__(self, index)
        return data, target, index

    return type(cls.__name__, (cls,), {
        '__getitem__': __getitem__,
    })
</code></pre>
<p>Then you can use it like so:</p>
<pre><code>MNISTWithIndices = dataset_with_indices(MNIST)
dataset = MNISTWithIndices('~/datasets/mnist')
</code></pre>
<p>This at least avoids the need to write a child class for every dataset you want to do this with.</p>

Yang_Zheng: <p>hi, i want to retrive all the indices of my dataset, could you give me some advice?</p>

ptrblck: <p>To get all indices, you could just create a range operation using <code>indices = range(len(dataset))</code>.<br>
If you want to get the current indices for the batch, <a href="https://discuss.pytorch.org/t/how-to-retrieve-the-sample-indices-of-a-mini-batch/7948/12">my code example</a> should work.<br>
Do you run into any issues with this code snippet?</p>