Batch size missing in data after torch.utils.data.random_split on dataloader.dataset

avtrulzz: <p>I used random_split() to divide my data into train and test and I observed that if random split is done after the dataloader is created, batch size is missing when getting a batch of data from the dataloader.</p>
<pre><code class="lang-auto">import torch
from torchvision import transforms, datasets
from torch.utils.data import random_split

# Normalize the data
transform_image = transforms.Compose([transforms.Resize((240, 320)),
                                    transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

data = '/data/imgs/train'

def load_dataset():
    data_path = data
    main_dataset = datasets.ImageFolder(
        root = data_path,
        transform = transform_image
    )

    loader = torch.utils.data.DataLoader(
        dataset = main_dataset,
        batch_size= 64,
        num_workers = 0,
        shuffle= True
    )

    # Dataset has 22424 data points
    trainloader, testloader = random_split(loader.dataset, [21000, 1424])

    return trainloader, testloader

trainloader, testloader = load_dataset()
</code></pre>
<p>Now to get a single batch of images from the train and test loaders:</p>
<pre><code class="lang-auto">images, labels = next(iter(trainloader))
images.shape
# %%
len(trainloader)

# %%
images_test, labels_test = next(iter(testloader))
images_test.shape

# %%
len(testloader)
</code></pre>
<p>The output that I get is does not have the batch size for train or test batches. Teh output dims should be [batch x channel x H x W] but I get [channel x H x W].</p>
<p>Output:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b621bdf7e831a35997ea6fc2b96f2c694fd90750.png" data-download-href="https://discuss.pytorch.org/uploads/default/b621bdf7e831a35997ea6fc2b96f2c694fd90750" title="Capture"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b621bdf7e831a35997ea6fc2b96f2c694fd90750.png" alt="Capture" data-base62-sha1="pZdbSK3p2gbF0Zn4iGWlFLo9PCU" width="541" height="500" data-dominant-color="F7F8F8"><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">Capture</span><span class="informations">961×888 14.4 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>
<p>But if I create the split from the dataset and then make two data loaders using the splits, I get the batchsize in the output.</p>
<pre><code class="lang-auto">def load_dataset():
    data_path = data
    main_dataset = datasets.ImageFolder(
        root = data_path,
        transform = transform_image
    )
    # Dataset has 22424 data points
    train_data, test_data = random_split(main_dataset, [21000, 1424])

    trainloader = torch.utils.data.DataLoader(
        dataset = train_data,
        batch_size= 64,
        num_workers = 0,
        shuffle= True
    )

    testloader = torch.utils.data.DataLoader(
        dataset = test_data,
        batch_size= 64,
        num_workers= 0,
        shuffle= True
    )

    return trainloader, testloader


trainloader, testloader = load_dataset()
</code></pre>
<p>On running the same 4 commands to get a single train and test batch:</p>
<pre><code class="lang-auto">images, labels = next(iter(trainloader))
images.shape
# %%
len(trainloader)

# %%
images_test, labels_test = next(iter(testloader))
images_test.shape

# %%
len(testloader)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; images, labels = next(iter(trainloader))
&gt;&gt;&gt; images.shape
torch.Size([64, 3, 240, 320])

&gt;&gt;&gt; len(trainloader)
329

&gt;&gt;&gt; images_test, labels_test = next(iter(testloader))
&gt;&gt;&gt; images_test.shape
torch.Size([64, 3, 240, 320])

len(testloader)
&gt;&gt; 23


</code></pre>
<p>Is the first approach wrong? Although the length shows that the data has been split. So why do I not see the batch size?</p>

ptrblck: <p><code>random_split</code> splits <code>Datasets</code> and returns <code>Datasets</code>, so <code>trainloader</code> and <code>testloader</code> should be of type <code>Dataset</code> not <code>DataLoader</code>.</p>
<p>The recommended way is your second approach: split the <code>Datasets</code> first, then wrap them in a <code>DataLoader</code>.</p>

avtrulzz: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>: Thanks for the explanation,  I amusing the second approach. But do you have any idea why the batch size would be missing in the first approach ?</p>

ptrblck: <p>Yes, because indexing a <code>Dataset</code> will return the sample without a batch dimension. The <code>DataLoader</code> creates the batch. Since <code>random_split</code> returns a <code>Dataset</code>, you won’t have a batch dimension using this approach.</p>

saba: <p>Hi Ptrblck,</p>
<p>I used thei command but dataloader works in correct, I think I should do sth before passing the index from split data. TrainData1 and  ValidationData1 are index.</p>
<pre><code class="lang-auto">    [TrainData1, ValidationData1]= train_test_split( np.arange(len(wholetargetArray)),train_size=0.7,test_size=0.3,stratify=wholetargetArray)
    trainloader=torch.utils.data.DataLoader(TrainData1, batch_size=64,shuffle=True,drop_last=True, num_workers=0)
    validationloader=torch.utils.data.DataLoader(ValidationData1, batch_size=6, drop_last=True,num_workers=0)```

what is your suggestion?

I appreciate your help</code></pre>

ptrblck: <p><code>TrainData1</code> and <code>ValidationData1</code> should contain the split indices, so you could pass these indices to a <code>torch.utils.data.Subset</code>.</p>