How to obtain the path name of train_dataset after using random_split in torch

user_16783344 0
Viewed 49

I have the following code:

import torch, torchvision
root_dataset ="./data"
dataset = torchvision.datasets.folder.ImageFolder(root=root_dataset, transform=None, target_transform=None)
train_dataset, valid_dataset = torch.utils.data.dataset.random_split(
    dataset=dataset,
    lengths=[num_train, num_valid]
)

My question is:

How can I obtain the name list of the path of train_dataset after using random_split in torch?

Thank you.

pytorch python
1 Answers

The paths (and labels) are stored in dataset.imgs. For instance, for imagenet:

In [ ]: print(dataset.imgs[0])
Out [ ]: ('/shareDB/imagenet/val/n01440764/ILSVRC2012_val_00000293.JPEG', 0)

After splitting the dataset, each split points to the original dataset:

In [ ]: len(train_dataset.dataset), len(valid_dataset.dataset)
Out [ ]: (50000, 50000)

However, each split also holds the indices of samples from the original dataset selected for the split. You can use these indices and the original dataset to get a list of the images selected for each split:

valid_imgs = [valid_dataset.dataset.imgs[i_] for i_ in valid_dataset.indices]
train_imgs = [train_dataset.dataset.imgs[i_] for i_ in train_dataset.indices]
user_17144100
Related