Как разделить тестовые и обучающие данные в наборе данных на основе количества целей в каждой категории - PullRequest
Новая
       6

Как разделить тестовые и обучающие данные в наборе данных на основе количества целей в каждой категории

0 голосов
/ 13 июля 2020

У меня есть imageFolder в PyTorch, который содержит мои категоризированные изображения данных. Каждая папка - это имя категории, а в папке - изображения этой категории.

Я загрузил данные и разделил данные поезда и теста через сэмплер со случайным train_test_split. Но проблема в том, что у меня плохое распределение данных, и в некоторых классах много изображений, а в некоторых меньше.

Итак, для решения этой проблемы я хочу выбрать 20% каждого класса в качестве тестовых данных и Остальное - данные поезда 

ds = ImageFolder(filePath, transform=transform)
batch_size = 64
validation_split = 0.2

indices = list(range(len(ds))) # indices of the dataset

# TODO: fix spliting
train_indices,test_indices = train_test_split(indices,test_size=0.2) 

# Creating PT data samplers and loaders:
train_sampler = SubsetRandomSampler(train_indices)
test_sampler = SubsetRandomSampler(test_indices)

train_loader = torch.utils.data.DataLoader(ds, batch_size=batch_size, sampler=train_sampler, num_workers=16)
test_loader = torch.utils.data.DataLoader(ds, batch_size=batch_size, sampler=test_sampler, num_workers=16)

Есть идеи, как это исправить?

Ответы [ 2 ]

2 голосов
/ 13 июля 2020

Используйте аргумент stratify в train_test_split в соответствии с docs . Если ваши индексы меток представляют собой массив, называемый y, выполните: 

train_indices,test_indices = train_test_split(indices, test_size=0.2, stratify=y)
0 голосов
/ 13 июля 2020

Попробуйте использовать StratifiedKFold или StratifiedShuffleSplit .

Согласно документам:

Этот объект перекрестной проверки является разновидностью KFold, возвращающий многослойные складки. Складки сделаны с сохранением процентного отношения образцов для каждого класса.

В вашем случае вы можете попробовать: 

from sklearn.model_selection import StratifiedShuffleSplit
sss = StratifiedShuffleSplit(n_splits=5, test_size=0.5, random_state=0)
for train_index, test_index in sss.split(ds):
    train = torch.utils.data.Subset(dataset, train_index)
    test = torch.utils.data.Subset(dataset, test_index)
    trainloader = torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=False)
    testloader = torch.utils.data.DataLoader(test, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=False)
...