/
datasets.py
157 lines (129 loc) · 5.43 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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import os
from typing import *
import torch
from torchvision import transforms, datasets
from torch.utils.data import Dataset
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 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
if os.path.exists("/home/bta/Workspace/data/imagenet"):
IMAGENET_LOC_ENV = "/home/bta/Workspace/data/imagenet/val/"
DATASET_LOC = '/home/bta/Workspace/data/imagenet/'
elif os.path.exists("/sensei-fs/users/qtran"):
IMAGENET_LOC_ENV = "/sensei-fs/users/qtran/imagenet/val/"
DATASET_LOC = '/sensei-fs/users/qtran/imagenet/'
elif os.path.exists("/home/quantran/AML/MOO_Ensemble/data"):
IMAGENET_LOC_ENV = "/home/quantran/AML/MOO_Ensemble/data"
DATASET_LOC = '/home/quantran/AML/MOO_Ensemble/data'
else:
print("Invalid Dataset Path, Exit")
exit()
os.environ[IMAGENET_LOC_ENV] = DATASET_LOC
# list of all datasets
DATASETS = ["imagenet", "cifar10", "mnist", "cifar100"]
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)
elif dataset == "mnist":
return _mnist(split)
elif dataset == 'cifar100':
return _cifar100(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
elif dataset == "mnist":
return 10
elif dataset == 'cifar100':
return 100
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)
elif dataset == "mnist":
return NormalizeLayer(_MNIST_MEAN, _MNIST_STDDEV)
elif dataset == 'cifar100':
return NormalizeLayer(_CIFAR100_MEAN, _CIFAR100_STDDEV)
_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]
_CIFAR100_MEAN = [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]
_CIFAR100_STDDEV = [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]
_MNIST_MEAN = [0.5, ]
_MNIST_STDDEV = [0.5, ]
def _mnist(split: str) -> Dataset:
if split == "train":
return datasets.MNIST(DATASET_LOC, train=True, download=True, transform=transforms.Compose([
transforms.RandomCrop(28),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
elif split == "test":
return datasets.MNIST(DATASET_LOC, train=False, transform=transforms.ToTensor())
def _cifar10(split: str) -> Dataset:
if split == "train":
return datasets.CIFAR10(DATASET_LOC, train=True, download=True, transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
elif split == "test":
return datasets.CIFAR10(DATASET_LOC, train=False, download=True, transform=transforms.ToTensor())
def _cifar100(split: str) -> Dataset:
if split == "train":
return datasets.CIFAR100(DATASET_LOC, train=True, download=True, transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
elif split == "test":
return datasets.CIFAR100(DATASET_LOC, 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.Resize(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).to(device)
self.sds = torch.tensor(sds).to(device)
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).to(input.device)
sds = self.sds.repeat((batch_size, height, width, 1)).permute(0, 3, 1, 2).to(input.device)
#print(input)
return (input - means) / sds