-
Notifications
You must be signed in to change notification settings - Fork 20
/
data_loader.py
54 lines (39 loc) · 1.78 KB
/
data_loader.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
import random
import numpy as np
from itertools import cycle
from utils import imshow_grid, transform_config
from torchvision import datasets
from torch.utils.data import Dataset, DataLoader
class MNIST_Paired(Dataset):
def __init__(self, root='mnist', download=True, train=True, transform=transform_config):
self.mnist = datasets.MNIST(root=root, download=download, train=train, transform=transform)
self.data_dict = {}
for i in range(self.__len__()):
image, label = self.mnist.__getitem__(i)
try:
self.data_dict[label.item()]
except KeyError:
self.data_dict[label.item()] = []
self.data_dict[label.item()].append(image)
def __len__(self):
return self.mnist.__len__()
def __getitem__(self, index):
image, label = self.mnist.__getitem__(index)
# return another image of the same class randomly selected from the data dictionary
# this is done to simulate pair-wise labeling of data
return image, random.SystemRandom().choice(self.data_dict[label.item()]), label
if __name__ == '__main__':
"""
test code for data loader
"""
mnist_paired = MNIST_Paired()
loader = cycle(DataLoader(mnist_paired, batch_size=16, shuffle=True, num_workers=0, drop_last=True))
print(mnist_paired.data_dict.keys())
image_batch, image_batch_2, labels_batch = next(loader)
print(labels_batch)
image_batch = np.transpose(image_batch, (0, 2, 3, 1))
image_batch = np.concatenate((image_batch, image_batch, image_batch), axis=3)
imshow_grid(image_batch)
image_batch_2 = np.transpose(image_batch_2, (0, 2, 3, 1))
image_batch_2 = np.concatenate((image_batch_2, image_batch_2, image_batch_2), axis=3)
imshow_grid(image_batch_2)