def get_custom_imagenet(restricted=False,
data_path='./data',
data_aug=False,
shuffle_val=False,
batch_size=20):
"""
We use helpers from robustness library
Returns:
loader: dataset loader
norm: normalization function for dataset
label_map: label map (class numbers to names) for dataset
"""
if restricted:
ds = RestrictedImageNet(data_path)
label_map = CLASS_DICT['RestrictedImageNet']
else:
ds = ImageNet(data_path)
label_map = CLASS_DICT['ImageNet']
label_map = {k: v.split(',')[0] for k, v in label_map.items()}
normalization = helpers.InputNormalize(ds.mean.cuda(), ds.std.cuda())
loaders = ds.make_loaders(1,
batch_size=batch_size,
data_aug=data_aug,
shuffle_val=shuffle_val)
return loaders, normalization, label_map
def __init__(self, model, dataset):
"""
Initialize the Attacker
Args:
nn.Module model : the PyTorch model to attack
Dataset dataset : dataset the model is trained on, only used to get mean and std for normalization
"""
super(AttackerFGSM, self).__init__()
self.normalize = helpers.InputNormalize(dataset.mean, dataset.std)
self.model = model
def load_dataset(dataset, batch_size, num_workers=1, data_path='./data'):
'''
Load pretrained model with specified architecture.
Args:
dataset (str): name of one of dataset
('restricted_imagenet' or 'imagenet')
batch_size (int): batch size
num_workers (int): number of workers
data_path (str): path to data
Returns:
ds: dataset object
loader: dataset loader
norm: normalization function for dataset
label_map: label map (class numbers to names) for dataset
'''
ds = DATASETS[dataset](data_path)
loaders = ds.make_loaders(num_workers, batch_size, data_aug=False)
normalization = helpers.InputNormalize(ds.mean, ds.std)
label_map = CLASS_DICT['ImageNet'] if dataset == 'imagenet' else CLASS_DICT[
'RestrictedImageNet']
return ds, loaders, normalization, label_map
def __init__(self, model, dataset):
super().__init__()
self.model = model
self.normalize = helpers.InputNormalize(dataset.mean, dataset.std)
CIFAR_LABEL_MAP = {
0: 'airplane',
1: 'automobile',
2: 'bird',
3: 'cat',
4: 'deer',
5: 'dog',
6: 'frog',
7: 'horse',
8: 'ship',
9: 'truck'
}
CIFAR_NORMALIZATION = helpers.InputNormalize(
torch.tensor([0.4914, 0.4822, 0.4465]).cuda(), # mean
torch.tensor([0.2023, 0.1994, 0.2010]).cuda() # std
)
class Subset(torch.utils.data.Dataset):
def __init__(self, dataset, targets, per_target, random=False):
"""
Picks subset of the dataset.
Arguments
random: if False, always picks the same indices for given parameters.
targets are converted to labels: [3,5,1] -> [0,1,2]
"""
self.dataset = dataset
self.indices = []
self.target_to_label = {t: i for i, t in enumerate(targets)}