CODEWITHSUNDEEP
machine-learningimage-processingcomputer-visionpytorch
toy_programmer
toy_programmer

Reputation: 101

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

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?

Upvotes: 9

Views: 11461

Answers (1)

Shai
Shai

Reputation: 114786

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

Upvotes: 14

Related Questions