Стратифицированное разделение поездов/валов/тестов в Pytorch ⇐ Python

[Anonymous]

У меня есть набор данных классификации изображений с 6 категориями, который я загружаю с помощью класса torchvision ImageFolder. Я написал ниже, чтобы разделить набор данных на 3 набора послойно:

Код: Выделить всё

from torch.utils.data import Subset
from sklearn.model_selection import train_test_split

train_indices, test_indices, _, _ = train_test_split
(
range(len(dataset)),
dataset.targets,
stratify=dataset.targets,
test_size=0.1,
random_state=1
)

train_dataset = Subset(dataset, train_indices)
test_dataset = Subset(dataset, test_indices)

train_targets = [label for _, label in train_dataset]
train_indices, val_indices, _, _ = train_test_split
(
range(len(train_dataset)),
train_targets,
stratify=train_targets,
test_size=0.111,
random_state=1
)

train_dataset = Subset(dataset, train_indices)

Это правильно разбивает данные на 90/10/10%. Все классы имеют одинаковое распределение между всеми тремя наборами, за исключением 6-го класса. В тестовую выборку попадают все образцы, относящиеся к 6 классу.
Что я делаю не так?

Подробнее здесь: https://stackoverflow.com/questions/730 ... in-pytorch