]) # test_preprocessor = transforms.ToTensor() dataset = args.dataset if dataset == "CIFAR-10": train_dataset = CIFAR10(IMAGE_DATA_ROOT[dataset], train=True, transform=train_preprocessor) # test_dataset = CIFAR10(IMAGE_DATA_ROOT[dataset], train=False, transform=test_preprocessor) elif dataset == "CIFAR-100": train_dataset = CIFAR100(IMAGE_DATA_ROOT[dataset], train=True, transform=train_preprocessor) # test_dataset = CIFAR100(IMAGE_DATA_ROOT[dataset], train=False, transform=test_preprocessor) elif dataset == "ImageNet": train_preprocessor = DataLoaderMaker.get_preprocessor(IMAGE_SIZE[dataset], True, center_crop=False) # test_preprocessor = DataLoaderMaker.get_preprocessor(IMAGE_SIZE[dataset], False, center_crop=True) train_dataset = ImageFolder(IMAGE_DATA_ROOT[dataset] + "/train", transform=train_preprocessor) elif dataset == "TinyImageNet": train_dataset = TinyImageNet(IMAGE_DATA_ROOT[dataset], train_preprocessor, train=True) batch_size = args.batch_size train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
def __init__(self, tot_num_tasks, dataset, inner_batch_size, protocol): """ Args: num_samples_per_class: num samples to generate "per class" in one batch batch_size: size of meta batch size (e.g. number of functions) """ self.img_size = IMAGE_SIZE[dataset] self.dataset = dataset if protocol == SPLIT_DATA_PROTOCOL.TRAIN_I_TEST_II: self.model_names = MODELS_TRAIN_STANDARD[self.dataset] elif protocol == SPLIT_DATA_PROTOCOL.TRAIN_II_TEST_I: self.model_names = MODELS_TEST_STANDARD[self.dataset] elif protocol == SPLIT_DATA_PROTOCOL.TRAIN_ALL_TEST_ALL: self.model_names = MODELS_TRAIN_STANDARD[ self.dataset] + MODELS_TEST_STANDARD[self.dataset] self.model_dict = {} for arch in self.model_names: if StandardModel.check_arch(arch, dataset): model = StandardModel(dataset, arch, no_grad=False).eval() if dataset != "ImageNet": model = model.cuda() self.model_dict[arch] = model is_train = True preprocessor = DataLoaderMaker.get_preprocessor( IMAGE_SIZE[dataset], is_train) if dataset == "CIFAR-10": train_dataset = CIFAR10(IMAGE_DATA_ROOT[dataset], train=is_train, transform=preprocessor) elif dataset == "CIFAR-100": train_dataset = CIFAR100(IMAGE_DATA_ROOT[dataset], train=is_train, transform=preprocessor) elif dataset == "MNIST": train_dataset = MNIST(IMAGE_DATA_ROOT[dataset], train=is_train, transform=preprocessor) elif dataset == "FashionMNIST": train_dataset = FashionMNIST(IMAGE_DATA_ROOT[dataset], train=is_train, transform=preprocessor) elif dataset == "TinyImageNet": train_dataset = TinyImageNet(IMAGE_DATA_ROOT[dataset], preprocessor, train=is_train) elif dataset == "ImageNet": preprocessor = DataLoaderMaker.get_preprocessor( IMAGE_SIZE[dataset], is_train, center_crop=True) sub_folder = "/train" if is_train else "/validation" # Note that ImageNet uses pretrainedmodels.utils.TransformImage to apply transformation train_dataset = ImageFolder(IMAGE_DATA_ROOT[dataset] + sub_folder, transform=preprocessor) self.train_dataset = train_dataset self.total_num_images = len(train_dataset) self.all_tasks = dict() all_images_indexes = np.arange(self.total_num_images).tolist() for i in range(tot_num_tasks): self.all_tasks[i] = { "image": random.sample(all_images_indexes, inner_batch_size), "arch": random.choice(list(self.model_dict.keys())) }