-
Notifications
You must be signed in to change notification settings - Fork 0
/
semidataset.py
211 lines (178 loc) · 8.17 KB
/
semidataset.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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
"""
Datasets with unlabeled (or pseudo-labeled) data
"""
from torchvision.datasets import CIFAR10, SVHN
from torch.utils.data import Sampler, Dataset
import torch
import numpy as np
import os
import pickle
import logging
DATASETS = ['cifar10', 'svhn']
class SemiSupervisedDataset(Dataset):
def __init__(self,
base_dataset='cifar10',
take_amount=None,
take_amount_seed=13,
add_svhn_extra=False,
aux_data_filename=None,
add_aux_labels=True,
aux_take_amount=None,
train=False,
**kwargs):
"""A dataset with auxiliary pseudo-labeled data"""
if base_dataset == 'cifar10':
self.dataset = CIFAR10(train=train, **kwargs)
elif base_dataset == 'svhn':
if train:
self.dataset = SVHN(split='train', **kwargs)
else:
self.dataset = SVHN(split='test', **kwargs)
# because torchvision is annoying
self.dataset.targets = self.dataset.labels
self.targets = list(self.targets)
if train and add_svhn_extra:
svhn_extra = SVHN(split='extra', **kwargs)
self.data = np.concatenate([self.data, svhn_extra.data])
self.targets.extend(svhn_extra.labels)
else:
raise ValueError('Dataset %s not supported' % base_dataset)
self.base_dataset = base_dataset
self.train = train
if self.train:
if take_amount is not None:
rng_state = np.random.get_state()
np.random.seed(take_amount_seed)
take_inds = np.random.choice(len(self.sup_indices),
take_amount, replace=False)
np.random.set_state(rng_state)
logger = logging.getLogger()
logger.info('Randomly taking only %d/%d examples from training'
' set, seed=%d, indices=%s',
take_amount, len(self.sup_indices),
take_amount_seed, take_inds)
self.targets = self.targets[take_inds]
self.data = self.data[take_inds]
self.sup_indices = list(range(len(self.targets)))
self.unsup_indices = []
if aux_data_filename is not None:
aux_path = os.path.join(kwargs['root'], aux_data_filename)
print("Loading data from %s" % aux_path)
with open(aux_path, 'rb') as f:
aux = pickle.load(f)
aux_data = aux['data']
aux_targets = aux['extrapolated_targets']
orig_len = len(self.data)
if aux_take_amount is not None:
rng_state = np.random.get_state()
np.random.seed(take_amount_seed)
take_inds = np.random.choice(len(aux_data),
aux_take_amount, replace=False)
np.random.set_state(rng_state)
logger = logging.getLogger()
logger.info(
'Randomly taking only %d/%d examples from aux data'
' set, seed=%d, indices=%s',
aux_take_amount, len(aux_data),
take_amount_seed, take_inds)
aux_data = aux_data[take_inds]
aux_targets = aux_targets[take_inds]
self.data = np.concatenate((self.data, aux_data), axis=0)
if not add_aux_labels:
self.targets.extend([-1] * len(aux_data))
else:
self.targets.extend(aux_targets)
# note that we use unsup indices to track the labeled datapoints
# whose labels are "fake"
self.unsup_indices.extend(
range(orig_len, orig_len+len(aux_data)))
logger = logging.getLogger()
logger.info("Training set")
logger.info("Number of training samples: %d", len(self.targets))
logger.info("Number of supervised samples: %d",
len(self.sup_indices))
logger.info("Number of unsup samples: %d", len(self.unsup_indices))
logger.info("Label (and pseudo-label) histogram: %s",
tuple(
zip(*np.unique(self.targets, return_counts=True))))
logger.info("Shape of training data: %s", np.shape(self.data))
# Test set
else:
self.sup_indices = list(range(len(self.targets)))
self.unsup_indices = []
logger = logging.getLogger()
logger.info("Test set")
logger.info("Number of samples: %d", len(self.targets))
logger.info("Label histogram: %s",
tuple(
zip(*np.unique(self.targets, return_counts=True))))
logger.info("Shape of data: %s", np.shape(self.data))
@property
def data(self):
return self.dataset.data
@data.setter
def data(self, value):
self.dataset.data = value
@property
def targets(self):
return self.dataset.targets
@targets.setter
def targets(self, value):
self.dataset.targets = value
def __len__(self):
return len(self.dataset)
def __getitem__(self, item):
self.dataset.labels = self.targets # because torchvision is annoying
return self.dataset[item]
def __repr__(self):
fmt_str = 'Semisupervised Dataset ' + self.__class__.__name__ + '\n'
fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
fmt_str += ' Training: {}\n'.format(self.train)
fmt_str += ' Root Location: {}\n'.format(self.dataset.root)
tmp = ' Transforms (if any): '
fmt_str += '{0}{1}\n'.format(tmp, self.dataset.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
tmp = ' Target Transforms (if any): '
fmt_str += '{0}{1}'.format(tmp, self.dataset.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
return fmt_str
class SemiSupervisedSampler(Sampler):
"""Balanced sampling from the labeled and unlabeled data"""
def __init__(self, sup_inds, unsup_inds, batch_size, unsup_fraction=0.5,
num_batches=None):
if unsup_fraction is None or unsup_fraction < 0:
self.sup_inds = sup_inds + unsup_inds
unsup_fraction = 0.0
else:
self.sup_inds = sup_inds
self.unsup_inds = unsup_inds
self.batch_size = batch_size
unsup_batch_size = int(batch_size * unsup_fraction)
self.sup_batch_size = batch_size - unsup_batch_size
if num_batches is not None:
self.num_batches = num_batches
else:
self.num_batches = int(
np.ceil(len(self.sup_inds) / self.sup_batch_size))
super().__init__(None)
def __iter__(self):
batch_counter = 0
while batch_counter < self.num_batches:
sup_inds_shuffled = [self.sup_inds[i]
for i in torch.randperm(len(self.sup_inds))]
for sup_k in range(0, len(self.sup_inds), self.sup_batch_size):
if batch_counter == self.num_batches:
break
batch = sup_inds_shuffled[sup_k:(sup_k + self.sup_batch_size)]
if self.sup_batch_size < self.batch_size:
batch.extend([self.unsup_inds[i] for i in
torch.randint(high=len(self.unsup_inds),
size=(
self.batch_size - len(
batch),),
dtype=torch.int64)])
# this shuffle operation is very important, without it
# batch-norm / DataParallel hell ensues
np.random.shuffle(batch)
yield batch
batch_counter += 1
def __len__(self):
return self.num_batches