Extended Dataset __getitem__ dimensions

christnp: <p>I’m trying on extend the PyTorch MNIST dataset to create a custom dataset of labeled pairs (so a tuple of (img1, img2, new_label)). I was thinking that the <code>__getitem__</code> function was the most appropriate entry point for this task. However, I’m getting weird results regarding the shape of the output. If I create a barebones extension (shown below, first snippet) I get the expected output shape of <code>torch.Size([10, 1, 28, 28])</code> for a batch size of 10. However, when I try access the data by index in the <code>__getitem__</code> function (second snippet below), it drops the first dimension of the image such that the output shape is <code>torch.Size([10, 28, 28])</code>. More importantly, when try to access the “third” dimension of the image in the <code>__getitem__</code> function I get an error because it does not exist. What am I doing wrong?</p>
<p>Snippet one, barebones extension of MNIST dataset (just passing the data):</p>
<pre><code class="lang-auto">class TestSubLoader(torchvision.datasets.MNIST):
    '''
    Subloader to extent MNIST dataset
    '''
    def __init__(self, *args, **kwargs):
        
        super(TestSubLoader, self).__init__(*args, **kwargs)
        
    def __len__(self):
        return len(self.data)
    
</code></pre>
<p>Which outputs:</p>
<pre><code class="lang-auto">__main__ sub_image: torch.Size([10, 1, 28, 28])
__main__ pri_image: torch.Size([10, 1, 28, 28])
</code></pre>
<p>Snippet two, adding in <code>__getitem__</code> to simply returning the batched data:</p>
<pre><code class="lang-auto">class TestSubLoader(torchvision.datasets.MNIST):
    '''
    Subloader to extent MNIST dataset
    '''
    def __init__(self, *args, **kwargs):
        
        super(TestSubLoader, self).__init__(*args, **kwargs)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx], self.targets[idx] # &lt;-- self.data[idx, :, :, :] = IndexError
</code></pre>
<p>Which outputs:</p>
<pre><code class="lang-auto">__main__ sub_image: torch.Size([10, 28, 28])
__main__ pri_image: torch.Size([10, 1, 28, 28])
</code></pre>
<p>Here is the code I use to run the test to get the outputs above:</p>
<pre><code class="lang-auto">mnist_set_subloaded = TestSubLoader(root=data_dir,train=True, 
                                   download=True,
                                   transform=transforms.Compose([transforms.ToTensor()])) 

mnist_set_primary = torchvision.datasets.MNIST(root=data_dir,train=True, 
                                   download=True,
                                   transform=transforms.Compose([transforms.ToTensor()]))
                                    

subloaded_loader = torch.utils.data.DataLoader(mnist_set_subloaded,
                                               shuffle=True,
                                               num_workers=2,
                                               batch_size=10)

primary_loader = torch.utils.data.DataLoader(mnist_set_primary,
                                               shuffle=True,
                                               num_workers=2,
                                               batch_size=10)

for i, ((sub_image,sub_target),(pri_image,pri_target)) \
        in enumerate(zip(subloaded_loader,primary_loader)):

    if i%10 == 0:
        print(f'__main__ sub_image: {sub_image.shape}') 
        print(f'__main__ pri_image: {pri_image.shape}') 
        break
</code></pre>
<p>First, is this a good approach to create image pairs (obviously, not implemented above…). Second, regarding the main topic, what am I missing – do I have to manually reshape the <code>self.data</code> to maintain the structure?</p>
<p>Thanks for your help.</p>

christnp: <p>This <a href="https://discuss.pytorch.org/t/did-i-misunderstand-data-dataset-getitem/53523">reference</a> helped me get in the right direction. It seems I need to convert to PIL and then be sure to include the transforms. So I have the code below. Is this accurate?</p>
<pre><code class="lang-auto">class TestSubLoader(torchvision.datasets.MNIST):
    '''
    Subloader to extent MNIST dataset
    '''
    def __init__(self, *args, **kwargs):
        
        super(TestSubLoader, self).__init__(*args, **kwargs)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        data, target = self.data[idx], self.targets[idx]
        # transform the data to PIL to get the right shape
        data = transforms.ToPILImage()(data)

        if self.transform is not None:
            data = self.transform(data)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return data, target
</code></pre>
<p>On to my other concern, is <code>__getitem__</code> a good place to collect image pairs? Any thoughts are welcome.</p>

ptrblck: <aside class="quote no-group" data-username="christnp" data-post="2" data-topic="106231">
<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/christnp/48/32366_2.png" class="avatar"> christnp:</div>
<blockquote>
<p>It seems I need to convert to PIL and then be sure to include the transforms.</p>
</blockquote>
</aside>
<p>Yes, this would basically recreate the original MNIST <code>__getitem__</code> method as seen <a href="https://github.com/pytorch/vision/blob/91e03b91fd9bab19b4c295692455a1883831a932/torchvision/datasets/mnist.py#L91-L111">here</a>.</p>
<p>Alternatively, you could also directly call into the <code>__getitem__</code> from the parent class via:</p>
<pre data-code-wrap="python"><code class="lang-python">class TestSubLoader(torchvision.datasets.MNIST):
    '''
    Subloader to extent MNIST dataset
    '''
    def __init__(self, *args, **kwargs):
        super(TestSubLoader, self).__init__(*args, **kwargs)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return super().__getitem__(idx)
</code></pre>
<p>which would yield the same output shapes as the original <code>MNIST</code> dataset.</p>
<aside class="quote no-group" data-username="christnp" data-post="2" data-topic="106231">
<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/christnp/48/32366_2.png" class="avatar"> christnp:</div>
<blockquote>
<p>On to my other concern, is <code>__getitem__</code> a good place to collect image pairs?</p>
</blockquote>
</aside>
<p>Yes, I think it’s the best place to add this logic to the data loading.<br>
Let me know, if it works for your use case or if you get stuck.</p>