How to split data into train and test sets using torchvision.datasets.Imagefolder?

Question

In my custom dataset, one kind of image is in one folder which torchvision.datasets.Imagefolder can handle, but how to split the dataset into train and test?

Does this answer your question? [How do I split a custom dataset into training and test datasets?](https://stackoverflow.com/questions/50544730/how-do-i-split-a-custom-dataset-into-training-and-test-datasets) — Aray Karjauv, May 24 '22 at 15:16

score 12 · Accepted Answer · answered Jul 29 '19 at 05:00

12

You can use torch.utils.data.Subset to split your ImageFolder dataset into train and test based on indices of the examples.
For example:

orig_set = torchvision.datasets.Imagefolder(...)  # your dataset
n = len(orig_set)  # total number of examples
n_test = int(0.1 * n)  # take ~10% for test
test_set = torch.utils.data.Subset(orig_set, range(n_test))  # take first 10%
train_set = torch.utils.data.Subset(orig_set, range(n_test, n))  # take the rest

answered Jul 29 '19 at 05:00

Shai

111,146
38
238
371

In this way, the rate between the label can be maintained? – toy_programmer Jul 29 '19 at 07:27
1

@toy_programmer you might need to be more careful I. the way you select indices for train/test – Shai Jul 29 '19 at 07:32
5

Agree with @Shai. If the data is ordered you will take only the first 10% of the classes – Enoon May 13 '20 at 09:18
better use [`random_split`](https://pytorch.org/docs/stable/data.html?highlight=random_split#torch.utils.data.random_split) – Aray Karjauv May 24 '22 at 15:18

How to split data into train and test sets using torchvision.datasets.Imagefolder?

1 Answers1

Linked