P.multi_gpu = False
### only use one ood_layer while training
P.ood_layer = P.ood_layer[0]
### Initialize dataset ###
train_set, test_set, image_size, n_classes = get_dataset(P, dataset=P.dataset)
P.image_size = image_size
P.n_classes = n_classes
if P.one_class_idx is not None:
cls_list = get_superclass_list(P.dataset)
P.n_superclasses = len(cls_list)
full_test_set = deepcopy(test_set) # test set of full classes
train_set = get_subclass_dataset(train_set, classes=cls_list[P.one_class_idx])
test_set = get_subclass_dataset(test_set, classes=cls_list[P.one_class_idx])
kwargs = {'pin_memory': False, 'num_workers': 4}
if P.multi_gpu:
train_sampler = DistributedSampler(train_set, num_replicas=P.n_gpus, rank=P.local_rank)
test_sampler = DistributedSampler(test_set, num_replicas=P.n_gpus, rank=P.local_rank)
train_loader = DataLoader(train_set, sampler=train_sampler, batch_size=P.batch_size, **kwargs)
test_loader = DataLoader(test_set, sampler=test_sampler, batch_size=P.test_batch_size, **kwargs)
else:
train_loader = DataLoader(train_set, shuffle=True, batch_size=P.batch_size, **kwargs)
test_loader = DataLoader(test_set, shuffle=False, batch_size=P.test_batch_size, **kwargs)
if P.ood_dataset is None:
if P.one_class_idx is not None:
transforms.Resize(256),
transforms.CenterCrop(256),
transforms.Resize(32),
transforms.ToTensor(),
])
# remove airliner(1), ambulance(2), parking_meter(18), schooner(22) since similar class exist in CIFAR-10
class_idx_list = list(range(30))
remove_idx_list = [1, 2, 18, 22]
for remove_idx in remove_idx_list:
class_idx_list.remove(remove_idx)
set_random_seed(0)
train_dir = os.path.join(IMAGENET_PATH, 'one_class_train')
Imagenet_set = datasets.ImageFolder(train_dir, transform=transform)
Imagenet_set = get_subclass_dataset(Imagenet_set, class_idx_list)
Imagenet_dataloader = DataLoader(Imagenet_set,
batch_size=100,
shuffle=True,
pin_memory=False)
total_test_image = None
for n, (test_image, target) in enumerate(Imagenet_dataloader):
if n == 0:
total_test_image = test_image
else:
total_test_image = torch.cat((total_test_image, test_image),
dim=0).cpu()
if total_test_image.size(0) >= 10000: