How to do a stratified split

username: Miguel_Campos
cooked: <p>Hello.</p>
<p>Sorry for my english, i am still learning and thanks you for help.</p>
<p>I have all my datas inside a torchvision.datasets.ImageFolder. The idea is split the data with stratified method. For that propoose, i am using torch.utils.data.SubsetRandomSampler of this way:</p>
<blockquote>
<p>dataset = torchvision.datasets.ImageFolder(train_dir, transform=train_transform)<br>
<strong>targets</strong> = dataset.targets</p>
</blockquote>
<p>Targets is a array of 0s and 1s (2-class classification) something like this:<br>
[0, 0, 1, 1, 0, 1,…]</p>
<blockquote>
<p>from sklearn.model_selection import <strong>train_test_split</strong></p>
<p><strong>train_idx, valid_idx</strong>= train_test_split(<br>
np.arange(len(<strong>targets</strong>)),<br>
test_size=0.2,<br>
shuffle=True,<br>
stratify=targets)</p>
</blockquote>
<p>I got stratified indices for train(80%) and valid(20%) something like this:</p>
<p>train_idx = [12, 54, 23, 123, 80…]<br>
valid_idx = [77, 09, 67, 122, 665…]</p>
<p>This indices indicates the position of my 0s and 1s in way that this are splitted in a balanced way.<br>
Then is when i try to split my real data generating my train/valid loader like this:</p>
<blockquote>
<p><strong>train_sampler</strong> = torch.utils.data.SubsetRandomSampler(<strong>train_idx</strong>)<br>
<strong>valid_sampler</strong> = torch.utils.data.SubsetRandomSampler(<strong>valid_idx</strong>)</p>
<p>train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=<strong>train_sampler</strong>)<br>
valid_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=<strong>valid_sampler</strong>)</p>
</blockquote>
<p>But len(train_loader.dataset) return whole dataset instead of 80% of them. What am i doing wrong?</p>

username: ptrblck
cooked: <p>You are not doing anything wrong and the batches provided by <code>train_loader</code> and <code>valid_loader</code> should just use the provided indices from the corresponding sampler.<br>
Since the underlying <code>.dataset</code> is not changed, the <code>len</code> will be the same as the original one.<br>
However, during the creation of the batches, the samplers will use the provided indices to load the data from the internal <code>Dataset</code>.</p>
<p>Alternatively, you could also wrap your <code>dataset</code> in a <code>Subset</code>, which will then change the length.<br>
Besides that both approaches should yield the same result.</p>

username: Miguel_Campos
cooked: <p>Oh! thanks you! i thought that the len in both .dataset should be different , but what you say makes sense. Thanks!</p>

username: ADONAI_TZEVAOT
cooked: <p>Hi, if I want to wrap the dataset in a Subset, how should I rewrite the code ? Please help.</p>

username: Amin_Nasim
cooked: <p>working with Subset is fairly easy. You pass the dataset and the indices to Subset. The benefit of using Subset over Sampler is that you can set shuffle = True in the DataLoader</p>
<pre><code class="lang-auto"># Stratified Sampling for train and val
train_idx, validation_idx = train_test_split(np.arange(len(train_data)),
                                             test_size=0.1,
                                             random_state=999,
                                             shuffle=True,
                                             stratify=train_data.img_labels)

# Subset dataset for train and val
train_dataset = Subset(train_data, train_idx)
validation_dataset = Subset(train_data, validation_idx)

# Dataloader for train and val
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
validation_loader = DataLoader(validation_dataset, batch_size=64, shuffle=False)
</code></pre>