-
Notifications
You must be signed in to change notification settings - Fork 181
/
helpers.py
197 lines (169 loc) · 6.27 KB
/
helpers.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
import torch as ch
import shutil
import dill
import os
from subprocess import Popen, PIPE
import pandas as pd
from PIL import Image
from . import constants
def has_attr(obj, k):
"""Checks both that obj.k exists and is not equal to None"""
try:
return (getattr(obj, k) is not None)
except KeyError as e:
return False
except AttributeError as e:
return False
def calc_est_grad(func, x, y, rad, num_samples):
B, *_ = x.shape
Q = num_samples//2
N = len(x.shape) - 1
with ch.no_grad():
# Q * B * C * H * W
extender = [1]*N
queries = x.repeat(Q, *extender)
noise = ch.randn_like(queries)
norm = noise.view(B*Q, -1).norm(dim=-1).view(B*Q, *extender)
noise = noise / norm
noise = ch.cat([-noise, noise])
queries = ch.cat([queries, queries])
y_shape = [1] * (len(y.shape) - 1)
l = func(queries + rad * noise, y.repeat(2*Q, *y_shape)).view(-1, *extender)
grad = (l.view(2*Q, B, *extender) * noise.view(2*Q, B, *noise.shape[1:])).mean(dim=0)
return grad
def ckpt_at_epoch(num):
return '%s_%s' % (num, constants.CKPT_NAME)
def accuracy(output, target, topk=(1,), exact=False):
"""
Computes the top-k accuracy for the specified values of k
Args:
output (ch.tensor) : model output (N, classes) or (N, attributes)
for sigmoid/multitask binary classification
target (ch.tensor) : correct labels (N,) [multiclass] or (N,
attributes) [multitask binary]
topk (tuple) : for each item "k" in this tuple, this method
will return the top-k accuracy
exact (bool) : whether to return aggregate statistics (if
False) or per-example correctness (if True)
Returns:
A list of top-k accuracies.
"""
with ch.no_grad():
# Binary Classification
if len(target.shape) > 1:
assert output.shape == target.shape, \
"Detected binary classification but output shape != target shape"
return [ch.round(ch.sigmoid(output)).eq(ch.round(target)).float().mean()], [-1.0]
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
res_exact = []
for k in topk:
correct_k = correct[:k].view(-1).float()
ck_sum = correct_k.sum(0, keepdim=True)
res.append(ck_sum.mul_(100.0 / batch_size))
res_exact.append(correct_k)
if not exact:
return res
else:
return res_exact
class InputNormalize(ch.nn.Module):
'''
A module (custom layer) for normalizing the input to have a fixed
mean and standard deviation (user-specified).
'''
def __init__(self, new_mean, new_std):
super(InputNormalize, self).__init__()
new_std = new_std[..., None, None]
new_mean = new_mean[..., None, None]
self.register_buffer("new_mean", new_mean)
self.register_buffer("new_std", new_std)
def forward(self, x):
x = ch.clamp(x, 0, 1)
x_normalized = (x - self.new_mean)/self.new_std
return x_normalized
class DataPrefetcher():
def __init__(self, loader, stop_after=None):
self.loader = loader
self.dataset = loader.dataset
self.stream = ch.cuda.Stream()
self.stop_after = stop_after
self.next_input = None
self.next_target = None
def __len__(self):
return len(self.loader)
def preload(self):
try:
self.next_input, self.next_target = next(self.loaditer)
except StopIteration:
self.next_input = None
self.next_target = None
return
with ch.cuda.stream(self.stream):
self.next_input = self.next_input.cuda(non_blocking=True)
self.next_target = self.next_target.cuda(non_blocking=True)
def __iter__(self):
count = 0
self.loaditer = iter(self.loader)
self.preload()
while self.next_input is not None:
ch.cuda.current_stream().wait_stream(self.stream)
input = self.next_input
target = self.next_target
self.preload()
count += 1
yield input, target
if type(self.stop_after) is int and (count > self.stop_after):
break
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
# ImageNet label mappings
def get_label_mapping(dataset_name, ranges):
if dataset_name == 'imagenet':
label_mapping = None
elif dataset_name == 'restricted_imagenet':
def label_mapping(classes, class_to_idx):
return restricted_label_mapping(classes, class_to_idx, ranges=ranges)
elif dataset_name == 'custom_imagenet':
def label_mapping(classes, class_to_idx):
return custom_label_mapping(classes, class_to_idx, ranges=ranges)
else:
raise ValueError('No such dataset_name %s' % dataset_name)
return label_mapping
def restricted_label_mapping(classes, class_to_idx, ranges):
range_sets = [
set(range(s, e+1)) for s,e in ranges
]
# add wildcard
# range_sets.append(set(range(0, 1002)))
mapping = {}
for class_name, idx in class_to_idx.items():
for new_idx, range_set in enumerate(range_sets):
if idx in range_set:
mapping[class_name] = new_idx
# assert class_name in mapping
filtered_classes = list(mapping.keys()).sort()
return filtered_classes, mapping
def custom_label_mapping(classes, class_to_idx, ranges):
mapping = {}
for class_name, idx in class_to_idx.items():
for new_idx, range_set in enumerate(ranges):
if idx in range_set:
mapping[class_name] = new_idx
filtered_classes = list(mapping.keys()).sort()
return filtered_classes, mapping