-
Notifications
You must be signed in to change notification settings - Fork 1
/
permute_masks.py
executable file
·111 lines (80 loc) · 4.87 KB
/
permute_masks.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
import numpy as np
from torchvision import transforms
from torchvision.datasets import CIFAR10, CIFAR100, FashionMNIST
from torch.utils.data import DataLoader, Subset
def cifar10_dataloaders(batch_size=128, data_dir = 'datasets/cifar10'):
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
])
train_set = Subset(CIFAR10(data_dir, train=True, transform=train_transform, download=True), list(range(45000)))
val_set = Subset(CIFAR10(data_dir, train=True, transform=test_transform, download=True), list(range(45000, 50000)))
test_set = CIFAR10(data_dir, train=False, transform=test_transform, download=True)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
return train_loader, val_loader, test_loader
def cifar10_dataloaders_index(batch_size=128, data_dir = 'datasets/cifar10', file_name=None):
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
])
#check index sequence
train_set_full = CIFAR10(data_dir, train=True, transform=train_transform, download=True)
all_labels = np.array(train_set_full.targets)
target_labels = np.load('npy_files/labels.npy')
equal_number = np.sum(all_labels==target_labels)
print('check index sequence {}'.format(equal_number))
assert equal_number == all_labels.shape[0]
#subset index
print('using {}'.format(file_name))
sub_index = np.load(file_name)
train_set = Subset(train_set_full, list(sub_index))
val_set = Subset(CIFAR10(data_dir, train=True, transform=test_transform, download=True), list(range(45000, 50000)))
test_set = CIFAR10(data_dir, train=False, transform=test_transform, download=True)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
return train_loader, val_loader, test_loader
def cifar100_dataloaders(batch_size=128, data_dir = 'datasets/cifar100'):
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
])
train_set = Subset(CIFAR100(data_dir, train=True, transform=train_transform, download=True), list(range(45000)))
val_set = Subset(CIFAR100(data_dir, train=True, transform=test_transform, download=True), list(range(45000, 50000)))
test_set = CIFAR100(data_dir, train=False, transform=test_transform, download=True)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
return train_loader, val_loader, test_loader
def fashionmnist_dataloaders(batch_size=64, data_dir = 'datasets/fashionmnist'):
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.Grayscale(3),
transforms.ToTensor(),
])
test_transform = transforms.Compose([
transforms.Grayscale(3),
transforms.ToTensor(),
])
train_set = Subset(FashionMNIST(data_dir, train=True, transform=train_transform, download=True), list(range(55000)))
val_set = Subset(FashionMNIST(data_dir, train=True, transform=test_transform, download=True), list(range(55000, 60000)))
test_set = FashionMNIST(data_dir, train=False, transform=test_transform, download=True)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
return train_loader, val_loader, test_loader