Как загрузить все 5 пакетов CIFAR10 в одной структуре данных, как в MNIST в PyTorch?

Question

С Mnist у меня есть один файл с метками и один файл для поезда, поэтому я просто делаю:

self.data = datasets.MNIST(root='./data', train=True, download=True)

В основном я создаю набор меток (от 0 до 9) и сохраняю i-я позиция изображения в структуре данных, для создания моих пользовательских задач:

   def make_tasks (self):
            self.task_to_examples = {} #task 0-9
            self.all_tasks = set (self.data.train_labels.numpy ())
            for i, digit in enumerate (self.data.train_labels.numpy ()):
                if str(digit) not in self.task_to_examples:
                    self.task_to_examples[str(digit)] = []
                self.task_to_examples[str(digits)].append(i)

Я не понимаю, как сделать то же самое, используя CIFAR10, потому что он разделен на 5 пакетов, я хотел бы все данные в одной структуре.

Answer 1

Если желаемой структурой является {"class_id": [indices of the samples]}, то для CIFAR10 вы можете сделать что-то вроде этого:

import numpy as np
import torchvision

# set root accordingly
cifar = torchvision.datasets.CIFAR10(root=".", train=True, download=True)
task_to_examples = {
    str(task_id): np.where(cifar.targets == task_id)[0].tolist()
    for task_id in np.unique(cifar.targets)
}