-
Notifications
You must be signed in to change notification settings - Fork 3
/
datasets.py
157 lines (128 loc) · 5.61 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
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"
# on aisecure gpu1
# os.environ[IMAGENET_LOC_ENV] = "/data/datasets/imagenet/ILSVRC2012"
# on asedl
# os.environ[IMAGENET_LOC_ENV] = "/srv/local/data/ImageNet/ILSVRC2012_full"
# on asedl but mount to aisecure gpu1
# os.environ[IMAGENET_LOC_ENV] = "/home/linyi2/data_mnt/imagenet/ILSVRC2012"
# on asedl but sync mount
# os.environ[IMAGENET_LOC_ENV] = "/home/linyi2/data_mnt/local_imagenet/data/ImageNet/ILSVRC2012_full"
# on aws server
# os.environ[IMAGENET_LOC_ENV] = '/data/imagenet/ILSVRC2012'
# on gpu3 server
os.environ[IMAGENET_LOC_ENV] = '/home/linyi/data/ILSVRC2012/'
# list of all datasets
DATASETS = ["imagenet", "cifar10", "mnist"]
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 == "fashionmnist":
return _fashion_mnist(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
def get_dataset_shape(dataset: str):
"""Return the number of classes in the dataset. """
if dataset == "imagenet":
return (3, 224, 224)
elif dataset == "cifar10":
return (3, 32, 32)
elif dataset == "mnist":
return (1, 28, 28)
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)
else:
raise Exception("Unknown dataset")
_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]
_MNIST_MEAN = [0.5]
_MNIST_STDDEV = [0.5]
_DEFAULT_MEAN = [0.5, 0.5, 0.5]
_DEFAULT_STDDEV = [0.5, 0.5, 0.5]
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)
def _mnist(split: str) -> Dataset:
if split == "train":
return datasets.MNIST("./dataset_cache", train=True, download=True, transform=transforms.ToTensor())
elif split == "test":
return datasets.MNIST("./dataset_cache", train=False, download=True, transform=transforms.ToTensor())
def _fashion_mnist(split: str) -> Dataset:
if split == "train":
return datasets.FashionMNIST("./dataset_cache", train=True, download=True, transform=transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
elif split == "test":
return datasets.FashionMNIST("./dataset_cache", train=False, download=True, transform=transforms.ToTensor())
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).contiguous()
sds = self.sds.repeat((batch_size, height, width, 1)).permute(0, 3, 1, 2).contiguous()
return (input - means) / sds