-
Notifications
You must be signed in to change notification settings - Fork 116
/
imbalance_cifar.py
189 lines (168 loc) · 7.94 KB
/
imbalance_cifar.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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
import torchvision
import torchvision.transforms as transforms
import os
import pickle
import numpy as np
class ImbalanceCIFAR10(torchvision.datasets.CIFAR10):
cls_num = 10
def __init__(self, root, imb_type='exp', imb_factor=0.01, rand_number=0, train=True,
transform=None, target_transform=None, download=False):
super(ImbalanceCIFAR10, self).__init__(root, train, transform, target_transform, download)
np.random.seed(rand_number)
img_num_list = self.get_img_num_per_cls(self.cls_num, imb_type, imb_factor)
self.gen_imbalanced_data(img_num_list)
def get_img_num_per_cls(self, cls_num, imb_type, imb_factor):
img_max = len(self.data) / cls_num
img_num_per_cls = []
if imb_type == 'exp':
for cls_idx in range(cls_num):
num = img_max * (imb_factor**(cls_idx / (cls_num - 1.0)))
img_num_per_cls.append(int(num))
elif imb_type == 'step':
for cls_idx in range(cls_num // 2):
img_num_per_cls.append(int(img_max))
for cls_idx in range(cls_num // 2):
img_num_per_cls.append(int(img_max * imb_factor))
else:
img_num_per_cls.extend([int(img_max)] * cls_num)
return img_num_per_cls
def gen_imbalanced_data(self, img_num_per_cls):
new_data = []
new_targets = []
targets_np = np.array(self.targets, dtype=np.int64)
classes = np.unique(targets_np)
# np.random.shuffle(classes)
self.num_per_cls_dict = dict()
for the_class, the_img_num in zip(classes, img_num_per_cls):
self.num_per_cls_dict[the_class] = the_img_num
idx = np.where(targets_np == the_class)[0]
np.random.shuffle(idx)
selec_idx = idx[:the_img_num]
new_data.append(self.data[selec_idx, ...])
new_targets.extend([the_class, ] * the_img_num)
new_data = np.vstack(new_data)
self.data = new_data
self.targets = new_targets
def get_cls_num_list(self):
cls_num_list = []
for i in range(self.cls_num):
cls_num_list.append(self.num_per_cls_dict[i])
return cls_num_list
class ImbalanceCIFAR100(ImbalanceCIFAR10):
base_folder = 'cifar-100-python'
url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
filename = "cifar-100-python.tar.gz"
tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
train_list = [
['train', '16019d7e3df5f24257cddd939b257f8d'],
]
test_list = [
['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
]
meta = {
'filename': 'meta',
'key': 'fine_label_names',
'md5': '7973b15100ade9c7d40fb424638fde48',
}
cls_num = 100
class SemiSupervisedImbalanceCIFAR10(torchvision.datasets.CIFAR10):
cls_num = 10
unlabel_size_factor = 5
def __init__(self, root, imb_type='exp', imb_factor=0.01, unlabel_imb_factor=1,
rand_number=0, train=True, transform=None, target_transform=None, download=False):
super(SemiSupervisedImbalanceCIFAR10, self).__init__(root, train, transform, target_transform, download)
# unlabeled
self.unlabeled_data = os.path.join(root, 'ti_80M_selected.pickle') # selected data from 80M-TI
self.unlabeled_pseudo = os.path.join(root, 'pseudo_labeled_cifar.pickle') # pseudo-label using model trained on imbalanced data
self.imb_factor = imb_factor
self.unlabel_imb_factor = unlabel_imb_factor
self.num_per_cls_dict = dict()
np.random.seed(rand_number)
img_num_list = self.get_img_num_per_cls(self.cls_num, imb_type, imb_factor)
img_num_list_unlabeled = self.get_img_num_per_cls_unlabeled(self.cls_num, img_num_list, unlabel_imb_factor)
self.gen_imbalanced_data(img_num_list, img_num_list_unlabeled)
def get_img_num_per_cls(self, cls_num, imb_type, imb_factor):
img_max = len(self.data) / cls_num
img_num_per_cls = []
if imb_type == 'exp':
for cls_idx in range(cls_num):
num = img_max * (imb_factor**(cls_idx / (cls_num - 1.0)))
img_num_per_cls.append(int(num))
elif imb_type == 'step':
for cls_idx in range(cls_num // 2):
img_num_per_cls.append(int(img_max))
for cls_idx in range(cls_num // 2):
img_num_per_cls.append(int(img_max * imb_factor))
else:
img_num_per_cls.extend([int(img_max)] * cls_num)
return img_num_per_cls
def get_img_num_per_cls_unlabeled(self, cls_num, labeled_img_num_list, imb_factor):
img_unlabeled_total = np.sum(labeled_img_num_list) * self.unlabel_size_factor
img_first_min = img_unlabeled_total // cls_num
img_num_per_cls_unlabel = []
for cls_idx in range(cls_num):
num = img_first_min * (imb_factor**(cls_idx / (cls_num - 1.0)))
img_num_per_cls_unlabel.append(int(num))
factor = img_unlabeled_total / np.sum(img_num_per_cls_unlabel)
img_num_per_cls_unlabel = [int(num * factor) for num in img_num_per_cls_unlabel]
print(f"Unlabeled est total:\t{img_unlabeled_total}\n"
f"After processing:\t{np.sum(img_num_per_cls_unlabel)},\t{img_num_per_cls_unlabel}")
return img_num_per_cls_unlabel
def gen_imbalanced_data(self, img_num_per_cls, img_num_per_cls_unlabeled):
new_data = []
new_targets = []
targets_np = np.array(self.targets, dtype=np.int64)
classes = np.unique(targets_np)
# np.random.shuffle(classes)
for the_class, the_img_num in zip(classes, img_num_per_cls):
self.num_per_cls_dict[the_class] = the_img_num
idx = np.where(targets_np == the_class)[0]
np.random.shuffle(idx)
selec_idx = idx[:the_img_num]
new_data.append(self.data[selec_idx, ...])
new_targets.extend([the_class, ] * the_img_num)
print(f"Labeled data extracted:\t{len(new_targets)}")
for i in range(self.cls_num):
print(self.num_per_cls_dict[i])
# unlabeled dataset
print("Loading unlabeled data from %s" % self.unlabeled_data)
with open(self.unlabeled_data, 'rb') as f:
aux = pickle.load(f)
aux_data = aux['data']
aux_truth = aux['extrapolated_targets']
print("Loading pseudo labels from %s" % self.unlabeled_pseudo)
with open(self.unlabeled_pseudo, 'rb') as f:
aux_targets = pickle.load(f)
aux_targets = aux_targets['extrapolated_targets']
for the_class, the_img_num in zip(classes, img_num_per_cls_unlabeled):
# ground truth is only used to select samples
idx = np.where(aux_truth == the_class)[0]
np.random.shuffle(idx)
selec_idx = idx[:the_img_num]
new_data.append(aux_data[selec_idx, ...])
# append pseudo-label
new_targets.extend(aux_targets[selec_idx])
for pseudo_class in aux_targets[selec_idx]:
self.num_per_cls_dict[pseudo_class] += 1
new_data = np.vstack(new_data)
self.data = new_data
self.targets = new_targets
assert new_data.shape[0] == len(new_targets), 'Length of data & labels do not match!'
print(f"Unlabeled data extracted:\t{len(new_targets)}")
for i in range(self.cls_num):
print(self.num_per_cls_dict[i])
def get_cls_num_list(self):
cls_num_list = []
for i in range(self.cls_num):
cls_num_list.append(self.num_per_cls_dict[i])
return cls_num_list
if __name__ == '__main__':
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)
trainset = SemiSupervisedImbalanceCIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = iter(trainset)
data, label = next(trainloader)
import pdb; pdb.set_trace()