def get_ImageNet(train=True, batch_size=None, shuffle=None, augm_type='test', num_workers=8, size=224, config_dict=None):
if batch_size == None:
if train:
batch_size = DEFAULT_TRAIN_BATCHSIZE
else:
batch_size = DEFAULT_TEST_BATCHSIZE
augm_config = {}
transform = get_imageNet_augmentation(type=augm_type, out_size=size, config_dict=augm_config)
if not train and augm_type != 'none':
print('Warning: ImageNet test set with ref_data augmentation')
if shuffle is None:
shuffle = train
path = get_imagenet_path()
if train == True:
dataset = datasets.ImageNet(path, split='train', transform=transform)
else:
dataset = datasets.ImageNet(path, split='val', transform=transform)
loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
shuffle=shuffle, num_workers=num_workers)
if config_dict is not None:
config_dict['Dataset'] = 'ImageNet'
config_dict['Batch out_size'] = batch_size
config_dict['Augmentation'] = augm_config
return loader
def get_imagenet(batch_size=64):
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
])
transform_test = transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize
])
train_loader = DataLoader(
datasets.ImageNet(root='.data/imagenet', train=True, download=False, transform=transform_train), batch_size=batch_size,
shuffle=True, drop_last=True
)
test_loader = DataLoader(
datasets.ImageNet(root='.data/imagenet', train=False, download=False, transform=transform_test), batch_size=batch_size*2,
shuffle=False, drop_last=True
)
return train_loader, test_loader
def load_data(self,
transform=None,
target_transform=None,
dataset='cifar10'):
if dataset == 'cifar10':
train_data = datasets.CIFAR10(self.data_dir,
train=True,
download=True,
transform=transform)
test_data = datasets.CIFAR10(self.data_dir,
train=False,
download=True,
transform=transform)
self.model_name += '_cifar'
elif dataset == 'vocseg':
train_data = datasets.VOCSegmentation(
self.data_dir,
image_set='train',
download=True,
transform=transform,
target_transform=target_transform)
test_data = datasets.VOCSegmentation(
self.data_dir,
image_set='val',
download=True,
transform=transform,
target_transform=target_transform)
self.model_name += '_vocseg'
elif dataset == 'imagenet':
train_data = datasets.ImageNet(self.data_dir,
train=True,
download=True,
transform=transform)
test_data = datasets.ImageNet(self.data_dir,
train=False,
download=True,
transform=transform)
self.model_name += '_imagenet'
else:
# raise an error?
pass
train_sampler, valid_sampler = self._train_valid_samplers(train_data)
self.train_loader = DataLoader(train_data,
batch_size=self.batch_size,
sampler=train_sampler,
num_workers=self.num_workers)
self.valid_loader = DataLoader(train_data,
batch_size=self.batch_size,
sampler=valid_sampler,
num_workers=self.num_workers)
self.test_loader = DataLoader(test_data,
batch_size=self.batch_size,
num_workers=self.num_workers)
def load_imagenet_datasets(val_ratio=0.2,
data_augmentation=False,
num_train_examples=None,
seed=42):
data_dir = os.path.join(os.path.dirname(__file__), "../data/imagenet/")
# add normalization
means = torch.tensor([0.485, 0.456, 0.406])
stds = torch.tensor([0.229, 0.224, 0.225])
normalize_transform = transforms.Normalize(mean=means, std=stds)
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize_transform,
])
test_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize_transform,
])
if not data_augmentation:
train_transform = test_transform
train_data = datasets.ImageNet(data_dir,
download=False,
split="train",
transform=train_transform)
val_data = datasets.ImageNet(data_dir,
download=False,
split="train",
transform=test_transform)
test_data = datasets.ImageNet(data_dir,
download=False,
split="val",
transform=test_transform)
# split train and validation
train_indices, val_indices = split(len(train_data), val_ratio, seed)
if num_train_examples is not None:
train_indices = np.random.choice(train_indices,
num_train_examples,
replace=False)
train_data = Subset(train_data, train_indices)
val_data = Subset(val_data, val_indices)
# name datasets and save statistics
for dataset in [train_data, val_data, test_data]:
dataset.dataset_name = "imagenet"
dataset.statistics = (means, stds)
return train_data, val_data, test_data, None
def main():
# args & device
args = config.get_args()
if torch.cuda.is_available():
print('Train on GPU!')
device = torch.device("cuda")
else:
device = torch.device("cpu")
# dataset
assert args.dataset in ['cifar10', 'imagenet']
train_transform, valid_transform = data_transforms(args)
if args.dataset == 'cifar10':
trainset = torchvision.datasets.CIFAR10(root=os.path.join(args.data_dir, 'cifar'), train=True,
download=True, transform=train_transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,
shuffle=True, pin_memory=True, num_workers=8)
valset = torchvision.datasets.CIFAR10(root=os.path.join(args.data_dir, 'cifar'), train=False,
download=True, transform=valid_transform)
val_loader = torch.utils.data.DataLoader(valset, batch_size=args.batch_size,
shuffle=False, pin_memory=True, num_workers=8)
elif args.dataset == 'imagenet':
train_data_set = datasets.ImageNet(os.path.join(args.data_dir, 'ILSVRC2012', 'train'), train_transform)
val_data_set = datasets.ImageNet(os.path.join(args.data_dir, 'ILSVRC2012', 'valid'), valid_transform)
train_loader = torch.utils.data.DataLoader(train_data_set, batch_size=args.batch_size, shuffle=True,
num_workers=8, pin_memory=True, sampler=None)
val_loader = torch.utils.data.DataLoader(val_data_set, batch_size=args.batch_size, shuffle=False,
num_workers=8, pin_memory=True)
# SinglePath_OneShot
choice = [3, 1, 2, 1, 0, 1, 3, 3, 1, 3, 0, 1, 0, 3, 3, 3, 3, 3, 0, 3]
#[2, 0, 2, 3, 2, 2, 3, 1, 2, 1, 0, 1, 0, 3, 1, 0, 0, 2, 3, 2]
model = SinglePath_Network(args.dataset, args.resize, args.classes, args.layers, choice)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, args.momentum, args.weight_decay)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda epoch: 1 - (epoch / args.epochs))
# flops & params & structure
flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32),) if args.dataset == 'cifar10'
else (torch.randn(1, 3, 224, 224),), verbose=False)
# print(model)
print('Random Path of the Supernet: Params: %.2fM, Flops:%.2fM' % ((params / 1e6), (flops / 1e6)))
model = model.to(device)
summary(model, (3, 32, 32) if args.dataset == 'cifar10' else (3, 224, 224))
# train supernet
start = time.time()
for epoch in range(args.epochs):
train(args, epoch, train_loader, device, model, criterion, optimizer, scheduler, supernet=False)
scheduler.step()
if (epoch + 1) % args.val_interval == 0:
validate(args, epoch, val_loader, device, model, criterion, supernet=False)
utils.save_checkpoint({'state_dict': model.state_dict(), }, epoch + 1, tag=args.exp_name)
utils.time_record(start)
def _download_data(self):
train_data = datasets.ImageNet(root="./data",
train=True,
download=True,
transform=self.transform_train)
test_data = datasets.ImageNet(root="./data",
train=False,
download=True,
transform=self.transform_test)
return train_data, test_data
def evaluate_acc(class_cls, criterion, args):
model = class_cls()
with original_state_dict_hooks(model):
model.load_state_dict(load_and_parse_state_dict(args.checkpoint),
strict=False)
model.cuda()
if args.spos_preprocessing:
train_trans = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4),
transforms.RandomHorizontalFlip(0.5),
ToBGRTensor()
])
else:
train_trans = transforms.Compose(
[transforms.RandomResizedCrop(224),
transforms.ToTensor()])
val_trans = transforms.Compose(
[transforms.RandomResizedCrop(224),
ToBGRTensor()])
train_dataset = datasets.ImageNet(args.imagenet_dir,
split='train',
transform=train_trans)
val_dataset = datasets.ImageNet(args.imagenet_dir,
split='val',
transform=val_trans)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.train_batch_size,
num_workers=args.workers,
shuffle=True)
test_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch_size,
num_workers=args.workers,
shuffle=True)
acc_before = test_acc(model, criterion, args.log_frequency, test_loader)
nni.report_intermediate_result(acc_before)
retrain_bn(model, criterion, args.train_iters, args.log_frequency,
train_loader)
acc = test_acc(model, criterion, args.log_frequency, test_loader)
assert isinstance(acc, float)
nni.report_intermediate_result(acc)
nni.report_final_result(acc)
def __init__(self,
data_dir,
batch_size,
shuffle=True,
validation_split=0.0,
num_workers=1,
training=True):
print('[INFO] Preparing the ImageNet dataset ...')
if training is True:
trans = transforms.Compose([
transforms.RandomCrop(224, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225)),
])
else:
trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225)),
])
self.data_dir = data_dir
if training is True:
split = 'train'
else:
split = 'val'
self.dataset = datasets.ImageNet(self.data_dir,
split=split,
download=True,
transform=trans)
super().__init__(self.dataset, batch_size, shuffle, validation_split,
num_workers)
def _get_org_dataset(self,
mode: str,
data_format: str = None,
**kwargs) -> datasets.DatasetFolder:
data_format = data_format or self.data_format
split = 'val' if mode == 'valid' else mode
return datasets.ImageNet(root=self.folder_path, split=split, **kwargs)
def __init__(self, path, split, transform=None, balanced = True):
self.imagenet = datasets.ImageNet(path, split=split, transform=transform)
self.balanced = balanced
self.num_classes = len(RESTRICTED_IMAGNET_RANGES)
class_idcs = []
for i in range(self.num_classes):
class_idcs.append([])
for i, label in enumerate(self.imagenet.targets):
for class_idx, (start, end) in enumerate(RESTRICTED_IMAGNET_RANGES):
if (label >= start) and (label <= end):
class_idcs[class_idx].append(i)
break
self.imagenet_linear_idcs = [item for sublist in class_idcs for item in sublist]
self.targets = [class_idx for class_idx in range(self.num_classes) for _ in class_idcs[class_idx]]
self.max_class_imgs = max([len(a) for a in class_idcs])
self.min_class_imgs = min([len(a) for a in class_idcs])
if not self.balanced:
self.length = sum([len(a) for a in class_idcs])
else:
self.length = self.num_classes * self.min_class_imgs
def main():
global best_acc
start_epoch = args.start_epoch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# Data loading
# traindir = os.path.join(args.data, 'train')
# valdir = os.path.join(args.data, 'val')
#normalize = transforms.Normalize()
'''
train_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(traindir, transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
),
batch_size=args.train_batch, shuffle=True,
num_workers=args.workers, pin_memory=True
)
'''
trainset = datasets.ImageNet('./imagenet_data', train=True, download=True)
trainloader = data.DataLoader(trainset, shuffle=True)
# create model
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
def get_data_loader(dataset_name, batch_size=64, shuffle=False):
"""
Returns a DataLoader with validation images for dataset_name
"""
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
])
if dataset_name == 'imagenet':
val_dataset = datasets.ImageNet(IMAGENET_VAL_DIR,
split='val',
transform=transform)
elif dataset_name == 'cifar10':
val_dataset = datasets.CIFAR10(CIFAR_VAL_DIR,
train=False,
transform=transform)
elif dataset_name == 'voc2012':
val_dataset = datasets.VOCDetection(VOC_VAL_DIR,
year='2012',
image_set='val',
transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=batch_size,
shuffle=shuffle,
num_workers=2)
return val_loader
def __init__(self, root, crop_size=96, data='voc2012'):
super().__init__()
## Root Directory
self.root = root
## Train Set
self.tr_data = data
if self.tr_data == 'stl-10':
self.data = datasets.STL10(self.root, split='train', download=True)
elif self.tr_data == 'imagenet':
self.data = datasets.ImageNet(self.root,
split='train',
download=True)
else:
self.data_list = []
for image in os.listdir(self.root):
h, w, _ = cv2.imread(self.root + '/' + image).shape
if h > 96 and w > 96:
self.data_list.append(self.root + '/' + image)
## Downscaling Factor
self.ds_factor = 4
## Transform functions
self.hr_transform = Compose([RandomCrop(crop_size, crop_size)])
self.lr_transform = Compose(
[Resize(width=crop_size // 4, height=crop_size // 4)])
def raw_dataset(self, data_dir: str, download: bool, train: bool,
transform):
split = ("train" if train else "val")
return datasets.ImageNet(data_dir,
download=download,
split=split,
transform=transform)
def __init__(self,
root,
train=True,
transform=None,
target_transform=None):
split = 'train' if train else 'val'
# using torchvision ImageNet to get classes
data_set = datasets.ImageNet(root,
split=split,
transform=transform,
target_transform=target_transform)
self.classes = list()
for class_tuple in data_set.classes:
self.classes.append(','.join(class_tuple))
self.length = len(data_set)
# get dataset
self.dbpath = os.path.join(root, f'{split}.lmdb')
# self.env = lmdb.open(self.dbpath, subdir=osp.isdir(self.dbpath),
# readonly=True, lock=False,
# readahead=False, meminit=False)
# with self.env.begin(write=False) as txn:
# self.length = loads_pyarrow(txn.get(b'__len__'))
# self.keys = loads_pyarrow(txn.get(b'__keys__'))
# get transform and target_transform
self.transform = transform
self.target_transform = target_transform
# create evaluator
self._update_evaluator()
def inf_train_gen_imagenet(batch_size, flip=True, train=True, infinity=True):
if flip:
transf = transforms.Compose([
transforms.CenterCrop(128),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(),
# transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])
else:
transf = transforms.ToTensor()
if train is True:
split = "train"
else:
split = "val"
loader = torch.utils.data.DataLoader(
datasets.ImageNet("/home/LargeData/ImageNet/",
split="train",
transform=transf),
batch_size,
drop_last=True,
shuffle=True,
num_workers=8,
)
if infinity is True:
while True:
for img, labels in loader:
yield img
else:
for img, labels in loader:
yield img
def load_imagenet(args):
# train_dir = os.path.join(args.dataset_path, 'train')
test_dir = os.path.join(args.dataset_path, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
# train_dataset = datasets.ImageFolder(
# train_dir,
# transforms.Compose([
# transforms.RandomResizedCrop(224),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# normalize,
# ]))
train_dataset = datasets.ImageNet(args.dataset_path,
split='train',
transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True,
sampler=None)
test_dataset = datasets.ImageNet(args.dataset_path,
split='val',
transform=transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.test_batch_size,
shuffle=False,
num_workers=4,
pin_memory=True)
return {'train': train_loader, 'test': test_loader}
def classify_train_dataset(data_dir,
transform=TinyImageNetTrainTransform,
**kwargs):
train_dir = os.path.join(data_dir, 'train')
return datasets.ImageNet(train_dir,
transform=transform,
download="False",
**kwargs)
def __init__(self, options):
self.dataset_root = abspath(
getattr_or_default(
options, 'dataset_root',
'/ssd_scratch/cvit/aditya.bharti/Imagenet-orig'))
# Repeat the other steps as for FewShotDataset. Maybe find a way to absorb this common code into something?
self.image_size = 256
self.image_channels = 3
self.mean = (0.5071, 0.4867, 0.4408)
self.std = (0.2675, 0.2565, 0.2761)
options.image_size = self.image_size
options.image_channels = self.image_channels
options.image_mean = self.mean
options.image_std = self.std
multi_transforms = get_transforms(options)
other_transform = tv_transforms.Compose([
tv_transforms.Resize(size=(self.image_size, self.image_size)),
tv_transforms.ToTensor(),
tv_transforms.Normalize(self.mean, self.std),
])
self.train_set = MultiTransformDataset(tv_datasets.ImageNet(
root=self.dataset_root,
split="train",
transform=tv_transforms.Resize(size=(self.image_size,
self.image_size))),
transform_list=multi_transforms,
options=options)
self.trainval_set = self.train_set
self.plain_train_set = tv_datasets.ImageNet(root=self.dataset_root,
split="train",
transform=other_transform)
self.plain_trainval_set = self.plain_train_set
self.test_set = tv_datasets.ImageNet(root=self.dataset_root,
split="val",
transform=other_transform)
self.valid_set = tv_datasets.ImageNet(root=self.dataset_root,
split="val",
transform=other_transform)
def classify_val_dataset(data_dir,
transform=TinyImageNetvalidationTransform,
**kwargs):
val_dir = os.path.join(data_dir, 'val')
return datasets.ImageNet(val_dir,
transform=transform,
download="False",
**kwargs)
def get_imagenet_labels():
path = get_imagenet_path()
dataset = datasets.ImageNet(path, split='val', transform='none')
classes_extended = dataset.classes
labels = []
for a in classes_extended:
labels.append(a[0])
return labels
def pre_benchmark_atk(**kwargs):
"""
Helper function that sets all the defaults while performing checks
for all the options passed before benchmarking attacks.
"""
# Set the Default options if nothing explicit provided
def_dict = {
'bs': 4,
'trf': get_trf('rz256_cc224_tt_normimgnet'),
'dset': 'NA',
'root': './',
'topk': (1, 5),
'dfunc': datasets.ImageFolder,
'download': True,
}
for key, val in def_dict.items():
if key not in kwargs: kwargs[key] = val
if kwargs['dset'] == 'NA':
if 'loader' not in kwargs:
dset = kwargs['dfunc'](kwargs['root'], transform=kwargs['trf'])
loader = DataLoader(dset, batch_size=kwargs['bs'], num_workers=2)
else:
loader = kwargs['loader']
# Set dataset specific functions here
else:
if kwargs['dset'] == IMGNET12:
dset = datasets.ImageNet(kwargs['root'],
split='test',
download=kwargs['download'],
transform=kwargs['trf'])
elif kwargs['dset'] == MNIST:
kwargs['trf'] = get_trf('tt_normmnist')
kwargs['dfunc'] = datasets.MNIST
dset = kwargs['dfunc'](kwargs['root'],
train=False,
download=kwargs['download'],
transform=kwargs['trf'])
else:
raise
loader = DataLoader(dset, shuffle=False, batch_size=kwargs['bs'])
topk = kwargs['topk']
for key, val in def_dict.items():
print('[INFO] Setting {} to {}.'.format(key, kwargs[key]))
# Deleting keys that is used just for benchmark_atk() function is
# important as the same kwargs dict is passed to initialize the attack
# So, otherwise the attack will throw an exception
for key in def_dict:
del kwargs[key]
if 'loader' in kwargs: del kwargs['loader']
return loader, topk, kwargs
def celebA_ImageNetOD(shuffle=True, batch_size=128, augm_type='default'):
augm = get_celebA_augmentation(augm_type)
root = get_imagenet_path()
dataset = datasets.ImageNet(root=root, split='train', transform=augm)
loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=shuffle,
num_workers=8)
return loader
def get_imagenet_label_wid_pairs():
path = get_imagenet_path()
dataset = datasets.ImageNet(path, split='val', transform='none')
classes_extended = dataset.classes
wids = dataset.wnids
label_wid_pairs = []
for a, b in zip(classes_extended, wids) :
label_wid_pairs.append((a[0], b))
return label_wid_pairs
def __init__(self):
self.train_transforms = transforms.Compose([
transforms.RandomResizedCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
self.test_transforms = transforms.Compose([
transforms.RandomResizedCrop(32),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
self.train_dset = datasets.ImageNet(root='./data/inet', train=True,
download=False, transform=self.train_transforms)
self.test_dset = datasets.ImageNet(root='./data/inet', train=False,
download=False, transform=self.test_transforms)
def raw_dataset(self, data_dir: str, download: bool, split: str,
transform):
assert split in ['train', 'val', 'test']
if split == 'val':
return None # no predetermined validation set
imagenet_split = ("train" if split == 'train' else "val")
# NOTE: as test set, the validation set of the ImageNet will be returned
return datasets.ImageNet(data_dir,
download=download,
split=imagenet_split,
transform=transform)
def initialize_folder(self):
try:
datasets.ImageNet(root=self.folder_path, split='train')
datasets.ImageNet(root=self.folder_path, split='val')
except RuntimeError:
raise RuntimeError(
'\n\n'
'You need to visit https://image-net.org/download-images.php '
'to download ImageNet.\n'
'There are direct links to files, but not legal to distribute. '
'Please apply for access permission and find links yourself.\n\n'
f'folder_path: {self.folder_path}\n'
'expected files:\n'
'{folder_path}/ILSVRC2012_devkit_t12.tar.gz\n'
'{folder_path}/ILSVRC2012_img_train.tar\n'
'{folder_path}/ILSVRC2012_img_val.tar\n'
'{folder_path}/meta.bin')
if not os.path.isdir(os.path.join(self.folder_path, 'valid')):
os.symlink(os.path.join(self.folder_path, 'val'),
os.path.join(self.folder_path, 'valid'))
def get_dataset(name='MNIST'):
if name == 'MNIST':
print("Loading MNIST")
dataset = datasets.MNIST(root='datasets/',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
num_channels = 1
elif name == 'Fashion MNIST':
print("Loading Fashion MNIST")
dataset = datasets.FashionMNIST(root='datasets/',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize((0.5, ),
(0.5, )),
]))
num_channels = 1
elif name == 'CIFAR 10':
print("Loading CIFAR10")
dataset = datasets.CIFAR10(root='datasets/',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
num_channels = 3
else:
print("Loading ImageNet")
dataset = datasets.ImageNet(root='datasets/',
train=True,
download=False,
transform=transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]))
num_channels = 3
return dataset, num_channels
def get_dataset(self):
"""
Uses torchvision.datasets.ImageNet to load dataset.
Downloads dataset if doesn't exist already.
Returns:
torch.utils.data.TensorDataset: trainset, valset
"""
trainset = datasets.ImageNet('datasets/ImageNet/train/',
split='train',
transform=self.train_transforms,
target_transform=None,
download=True)
valset = datasets.ImageNet('datasets/ImageNet/val/',
split='val',
transform=self.val_transforms,
target_transform=None,
download=True)
return trainset, valset
def get_imagenetWID_complement(wids):
path = get_imagenet_path()
imagenet = datasets.ImageNet(path, split='val', transform='none')
imagenet_wids = imagenet.wnids
complement = []
for wid in imagenet_wids:
if wid not in wids:
complement.append(wid)
return complement
