-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathDatasets.py
More file actions
140 lines (106 loc) · 5.86 KB
/
Datasets.py
File metadata and controls
140 lines (106 loc) · 5.86 KB
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
import random
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from Utils import seed_worker
class Dataset:
def __init__(self, root, batch_size, train_set_size, val_set_size, *, seed=None):
self.root = root
self.batch_size = batch_size
self.train_set_size = train_set_size
self.val_set_size = val_set_size
self.transform = transforms.Compose(
[transforms.ToTensor()]
)
self.g = torch.Generator()
if seed is not None:
self.seed_everything(seed)
def seed_everything(self, seed):
# https://pytorch.org/docs/stable/notes/randomness.html
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
# https://pytorch.org/docs/stable/notes/randomness.html#dataloader
self.g.manual_seed(seed)
def get_train_loader(self):
return self.train_loader
def get_val_loader(self):
return self.val_loader
def get_test_loader(self):
return self.test_loader
def get_labels(self):
return self.labels
class MNIST(Dataset):
def __init__(self, root, batch_size, train_set_size, val_set_size, *, seed=None):
super().__init__(root, batch_size, train_set_size, val_set_size, seed=seed)
# Get datasets
train_dataset = torchvision.datasets.MNIST(
root=self.root, train=True, download=True, transform=self.transform)
test_dataset = torchvision.datasets.MNIST(
root=self.root, train=False, download=True, transform=self.transform)
# Split train and validation sets
train_set, val_set = torch.utils.data.random_split(
train_dataset, [self.train_set_size, self.val_set_size])
# Load train, validation, and test sets
self.train_loader = torch.utils.data.DataLoader(
train_set, batch_size=self.batch_size, shuffle=True, worker_init_fn=seed_worker, generator=self.g)
self.val_loader = torch.utils.data.DataLoader(val_set, batch_size=self.batch_size, shuffle=False)
self.test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=self.batch_size, shuffle=True)
# Index to label
self.labels = [str(x) for x in range(10)]
class FashionMNIST(Dataset):
def __init__(self, root, batch_size, train_set_size, val_set_size, *, seed=None):
super().__init__(root, batch_size, train_set_size, val_set_size, seed=seed)
# Get datasets
train_dataset = torchvision.datasets.FashionMNIST(
root=self.root, train=True, download=True, transform=self.transform)
test_dataset = torchvision.datasets.FashionMNIST(
root=self.root, train=False, download=True, transform=self.transform)
# Split train and validation sets
train_set, val_set = torch.utils.data.random_split(
train_dataset, [self.train_set_size, self.val_set_size])
# Load train, validation, and test sets
self.train_loader = torch.utils.data.DataLoader(
train_set, batch_size=self.batch_size, shuffle=True, worker_init_fn=seed_worker, generator=self.g)
self.val_loader = torch.utils.data.DataLoader(val_set, batch_size=self.batch_size, shuffle=False)
self.test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=self.batch_size, shuffle=True)
# Index to label
self.labels = [
"T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
class Omniglot(Dataset):
def __init__(self, root, batch_size, train_set_size, val_set_size, *, seed=None):
super().__init__(root, batch_size, train_set_size, val_set_size, seed=seed)
# Get datasets
train_dataset = torchvision.datasets.Omniglot(
root=self.root, background=True, download=True, transform=self.transform)
test_dataset = torchvision.datasets.Omniglot(
root=self.root, background=False, download=True, transform=self.transform)
# Split train and validation sets
train_set, val_set = torch.utils.data.random_split(
train_dataset, [self.train_set_size, self.val_set_size])
# Load train, validation, and test sets
self.train_loader = torch.utils.data.DataLoader(
train_set, batch_size=self.batch_size, shuffle=True, worker_init_fn=seed_worker, generator=self.g)
self.val_loader = torch.utils.data.DataLoader(val_set, batch_size=self.batch_size, shuffle=False)
self.test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=self.batch_size, shuffle=True)
# Index to label
self.labels = [str(x) for x in range(1623)]
class CIFAR10(Dataset):
def __init__(self, root, batch_size, train_set_size, val_set_size, *, seed=None):
super().__init__(root, batch_size, train_set_size, val_set_size, seed=seed)
# Get datasets
train_dataset = torchvision.datasets.CIFAR10(
root=self.root, train=True, download=True, transform=self.transform)
test_dataset = torchvision.datasets.CIFAR10(
root=self.root, train=False, download=True, transform=self.transform)
# Split train and validation sets
train_set, val_set = torch.utils.data.random_split(
train_dataset, [self.train_set_size, self.val_set_size])
# Load train, validation, and test sets
self.train_loader = torch.utils.data.DataLoader(
train_set, batch_size=self.batch_size, shuffle=True, worker_init_fn=seed_worker, generator=self.g)
self.val_loader = torch.utils.data.DataLoader(val_set, batch_size=self.batch_size, shuffle=False)
self.test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=self.batch_size, shuffle=True)
# Index to label
self.labels = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]