-
Notifications
You must be signed in to change notification settings - Fork 2
/
datasets.py
135 lines (110 loc) · 5.15 KB
/
datasets.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
'''
- this file is necessary to load the datasets (CIFAR10 and ImageNet)
- it is based on the publicly available code https://github.com/locuslab/smoothing/blob/master/code/datasets.py written by Jeremy Cohen and https://github.com/Hadisalman/smoothing-adversarial/blob/master/code/datasets.py written by Hadi Salman
'''
from torchvision import transforms, datasets
from typing import *
import torch
import os
from torch.utils.data import Dataset
# set this environment variable to the location of your imagenet directory if you want to read ImageNet data.
# make sure your val directory is preprocessed to look like the train directory, e.g. by running this script
# https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh
IMAGENET_LOC_ENV = "IMAGENET_DIR"
# list of all datasets
DATASETS = ["imagenet", "cifar10"]
def get_dataset(dataset: str, split: str) -> Dataset:
"""Return the dataset as a PyTorch Dataset object"""
if dataset == "imagenet":
return _imagenet(split)
elif dataset == "cifar10":
return _cifar10(split)
def get_num_classes(dataset: str):
"""Return the number of classes in the dataset. """
if dataset == "imagenet":
return 1000
elif dataset == "cifar10":
return 10
def get_normalize_layer(dataset: str) -> torch.nn.Module:
"""Return the dataset's normalization layer"""
if dataset == "imagenet":
return NormalizeLayer(_IMAGENET_MEAN, _IMAGENET_STDDEV)
elif dataset == "cifar10":
return NormalizeLayer(_CIFAR10_MEAN, _CIFAR10_STDDEV)
# added for some SmoothAdv models
def get_input_center_layer(dataset: str) -> torch.nn.Module:
"""Return the dataset's Input Centering layer"""
if dataset == "imagenet":
return InputCenterLayer(_IMAGENET_MEAN)
elif dataset == "cifar10":
return InputCenterLayer(_CIFAR10_MEAN)
# added for some SmoothAdv models
class InputCenterLayer(torch.nn.Module):
"""Centers the channels of a batch of images by subtracting the dataset mean.
In order to certify radii in original coordinates rather than standardized coordinates, we
add the Gaussian noise _before_ standardizing, which is why we have standardization be the first
layer of the classifier rather than as a part of preprocessing as is typical.
"""
def __init__(self, means: List[float]):
"""
:param means: the channel means
:param sds: the channel standard deviations
"""
super(InputCenterLayer, self).__init__()
self.means = torch.tensor(means).cuda()
def forward(self, input: torch.tensor):
(batch_size, num_channels, height, width) = input.shape
means = self.means.repeat((batch_size, height, width, 1)).permute(0, 3, 1, 2)
return input - means
_IMAGENET_MEAN = [0.485, 0.456, 0.406]
_IMAGENET_STDDEV = [0.229, 0.224, 0.225]
_CIFAR10_MEAN = [0.4914, 0.4822, 0.4465]
_CIFAR10_STDDEV = [0.2023, 0.1994, 0.2010]
def _cifar10(split: str) -> Dataset:
if split == "train":
return datasets.CIFAR10("./dataset_cache", train=True, download=True, transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
elif split == "test":
return datasets.CIFAR10("./dataset_cache", train=False, download=True, transform=transforms.ToTensor())
def _imagenet(split: str) -> Dataset:
if not IMAGENET_LOC_ENV in os.environ:
raise RuntimeError("environment variable for ImageNet directory not set")
dir = os.environ[IMAGENET_LOC_ENV]
if split == "train":
subdir = os.path.join(dir, "train")
transform = transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
])
elif split == "test":
subdir = os.path.join(dir, "val")
transform = transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor()
])
return datasets.ImageFolder(subdir, transform)
class NormalizeLayer(torch.nn.Module):
"""Standardize the channels of a batch of images by subtracting the dataset mean
and dividing by the dataset standard deviation.
In order to certify radii in original coordinates rather than standardized coordinates, we
add the Gaussian noise _before_ standardizing, which is why we have standardization be the first
layer of the classifier rather than as a part of preprocessing as is typical.
"""
def __init__(self, means: List[float], sds: List[float]):
"""
:param means: the channel means
:param sds: the channel standard deviations
"""
super(NormalizeLayer, self).__init__()
self.means = torch.tensor(means).cuda()
self.sds = torch.tensor(sds).cuda()
def forward(self, input: torch.tensor):
(batch_size, num_channels, height, width) = input.shape
means = self.means.repeat((batch_size, height, width, 1)).permute(0, 3, 1, 2)
sds = self.sds.repeat((batch_size, height, width, 1)).permute(0, 3, 1, 2)
return (input - means) / sds