-
Notifications
You must be signed in to change notification settings - Fork 63
/
Copy pathdatasets.py
113 lines (83 loc) · 2.75 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
from pathlib import Path
import torch
import torch.nn.functional as F
from torchvision import transforms, datasets
n_bits = 8
def preprocess(x):
# Follows:
# https://github.com/tensorflow/tensor2tensor/blob/e48cf23c505565fd63378286d9722a1632f4bef7/tensor2tensor/models/research/glow.py#L78
x = x * 255 # undo ToTensor scaling to [0,1]
n_bins = 2 ** n_bits
if n_bits < 8:
x = torch.floor(x / 2 ** (8 - n_bits))
x = x / n_bins - 0.5
return x
def postprocess(x):
x = torch.clamp(x, -0.5, 0.5)
x += 0.5
x = x * 2 ** n_bits
return torch.clamp(x, 0, 255).byte()
def one_hot_encode(target):
"""
One hot encode with fixed 10 classes
Args: target - the target labels to one-hot encode
Retn: one_hot_encoding - the OHE of this tensor
"""
num_classes = 10
one_hot_encoding = F.one_hot(torch.tensor(target),num_classes)
return one_hot_encoding
def get_CIFAR10(augment, dataroot, download):
image_shape = (32, 32, 3)
num_classes = 10
if augment:
transformations = [
transforms.RandomAffine(0, translate=(0.1, 0.1)),
transforms.RandomHorizontalFlip(),
]
else:
transformations = []
transformations.extend([transforms.ToTensor(), preprocess])
train_transform = transforms.Compose(transformations)
test_transform = transforms.Compose([transforms.ToTensor(), preprocess])
path = Path(dataroot) / "data" / "CIFAR10"
train_dataset = datasets.CIFAR10(
path,
train=True,
transform=train_transform,
target_transform=one_hot_encode,
download=download,
)
test_dataset = datasets.CIFAR10(
path,
train=False,
transform=test_transform,
target_transform=one_hot_encode,
download=download,
)
return image_shape, num_classes, train_dataset, test_dataset
def get_SVHN(augment, dataroot, download):
image_shape = (32, 32, 3)
num_classes = 10
if augment:
transformations = [transforms.RandomAffine(0, translate=(0.1, 0.1))]
else:
transformations = []
transformations.extend([transforms.ToTensor(), preprocess])
train_transform = transforms.Compose(transformations)
test_transform = transforms.Compose([transforms.ToTensor(), preprocess])
path = Path(dataroot) / "data" / "SVHN"
train_dataset = datasets.SVHN(
path,
split="train",
transform=train_transform,
target_transform=one_hot_encode,
download=download,
)
test_dataset = datasets.SVHN(
path,
split="test",
transform=test_transform,
target_transform=one_hot_encode,
download=download,
)
return image_shape, num_classes, train_dataset, test_dataset