-
Notifications
You must be signed in to change notification settings - Fork 2
/
smooth.py
190 lines (162 loc) · 8.1 KB
/
smooth.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
import torch
from scipy.stats import norm, binom_test
import numpy as np
from math import ceil
from statsmodels.stats.proportion import proportion_confint
from distribution import Distribution
def sample_noise(model: torch.nn.Module, x: torch.tensor, dist: Distribution, label: int, num: int, num_classes: int, batch_size: int) -> np.ndarray:
""" Sample the base classifier's prediction under noisy corruptions of the input x.
:param x: the input [channel x width x height]
:param num: number of samples to collect
:param batch_size:
:return: votes for the true class
"""
tot = num
with torch.no_grad():
# counts = np.zeros(num_classes, dtype=int)
counts = 0
for i in range(ceil(num / batch_size)):
print('batch', i, 'num', num, end='\r', flush=True)
this_batch_size = min(batch_size, num)
num -= this_batch_size
batch = x.repeat((this_batch_size, 1, 1, 1))
# noise = dist.sample(this_batch_size).astype(np.float32)
noise = dist.sample(this_batch_size, cuda=True)
noise = torch.tensor(noise, device='cuda').resize_as(batch)
predictions = model(batch + noise).argmax(1)
counts += np.sum(predictions.cpu().numpy() == label)
# for idx in predictions.cpu().numpy():
# counts[idx] += 1
if counts <= 1000 and tot - num >= 4000:
# early halt
break
return counts, tot - num
def full_sample_noise(model: torch.nn.Module, x: torch.tensor, dist: Distribution, num: int, num_classes: int, batch_size: int) -> np.ndarray:
""" Sample the base classifier's prediction under noisy corruptions of the input x.
:param x: the input [channel x width x height]
:param num: number of samples to collect
:param batch_size:
:return: an ndarray[int] of length num_classes containing the per-class counts
"""
tot = num
counts = None
with torch.no_grad():
# counts = np.zeros(num_classes, dtype=int)
counts = 0
for i in range(ceil(num / batch_size)):
# print('batch', i, 'num', num, end='\r', flush=True)
this_batch_size = min(batch_size, num)
num -= this_batch_size
batch = x.repeat((this_batch_size, 1, 1, 1))
# noise = dist.sample(this_batch_size).astype(np.float32)
noise = dist.sample(this_batch_size, cuda=True)
noise = torch.tensor(noise, device='cuda').resize_as(batch)
logits = model(batch + noise)
predictions = torch.nn.functional.one_hot(logits.argmax(1), num_classes=logits.shape[-1])
predictions = predictions.sum(dim=0)
if counts is None:
counts = predictions
else:
counts += predictions
return counts
def confidence_bound(NA: int, N: int, alpha: float) -> (float, float):
""" Returns a (1 - alpha) confidence *interval* on a bernoulli proportion.
This function uses the Clopper-Pearson method.
:param NA: the number of "successes"
:param N: the number of total draws
:param alpha: the confidence level
:return: a lower bound on the binomial proportion which holds true w.p at least (1 - alpha) over the samples
"""
ret = proportion_confint(NA, N, alpha=alpha, method="beta")
return ret[0], ret[1]
class Smooth(object):
"""A smoothed classifier g """
# to abstain, Smooth returns this int
ABSTAIN = -1
def __init__(self, base_classifier: torch.nn.Module, num_classes: int, sigma: float):
"""
:param base_classifier: maps from [batch x channel x height x width] to [batch x num_classes]
:param num_classes:
:param sigma: the noise level hyperparameter
"""
self.base_classifier = base_classifier
self.num_classes = num_classes
self.sigma = sigma
def certify(self, x: torch.tensor, n0: int, n: int, alpha: float, batch_size: int) -> (int, float):
""" Monte Carlo algorithm for certifying that g's prediction around x is constant within some L2 radius.
With probability at least 1 - alpha, the class returned by this method will equal g(x), and g's prediction will
robust within a L2 ball of radius R around x.
:param x: the input [channel x height x width]
:param n0: the number of Monte Carlo samples to use for selection
:param n: the number of Monte Carlo samples to use for estimation
:param alpha: the failure probability
:param batch_size: batch size to use when evaluating the base classifier
:return: (predicted class, certified radius)
in the case of abstention, the class will be ABSTAIN and the radius 0.
"""
self.base_classifier.eval()
# draw samples of f(x+ epsilon)
counts_selection = self._sample_noise(x, n0, batch_size)
# use these samples to take a guess at the top class
cAHat = counts_selection.argmax().item()
# draw more samples of f(x + epsilon)
counts_estimation = self._sample_noise(x, n, batch_size)
# use these samples to estimate a lower bound on pA
nA = counts_estimation[cAHat].item()
pABar = self._lower_confidence_bound(nA, n, alpha)
if pABar < 0.5:
return Smooth.ABSTAIN, 0.0
else:
radius = self.sigma * norm.ppf(pABar)
return cAHat, radius
def predict(self, x: torch.tensor, n: int, alpha: float, batch_size: int) -> int:
""" Monte Carlo algorithm for evaluating the prediction of g at x. With probability at least 1 - alpha, the
class returned by this method will equal g(x).
This function uses the hypothesis test described in https://arxiv.org/abs/1610.03944
for identifying the top category of a multinomial distribution.
:param x: the input [channel x height x width]
:param n: the number of Monte Carlo samples to use
:param alpha: the failure probability
:param batch_size: batch size to use when evaluating the base classifier
:return: the predicted class, or ABSTAIN
"""
self.base_classifier.eval()
counts = self._sample_noise(x, n, batch_size)
top2 = counts.argsort()[::-1][:2]
count1 = counts[top2[0]]
count2 = counts[top2[1]]
if binom_test(count1, count1 + count2, p=0.5) > alpha:
return Smooth.ABSTAIN
else:
return top2[0]
def _sample_noise(self, x: torch.tensor, num: int, batch_size) -> np.ndarray:
""" Sample the base classifier's prediction under noisy corruptions of the input x.
:param x: the input [channel x width x height]
:param num: number of samples to collect
:param batch_size:
:return: an ndarray[int] of length num_classes containing the per-class counts
"""
with torch.no_grad():
counts = np.zeros(self.num_classes, dtype=int)
for _ in range(ceil(num / batch_size)):
this_batch_size = min(batch_size, num)
num -= this_batch_size
batch = x.repeat((this_batch_size, 1, 1, 1))
noise = torch.randn_like(batch, device='cuda') * self.sigma
predictions = self.base_classifier(batch + noise).argmax(1)
counts += self._count_arr(predictions.cpu().numpy(), self.num_classes)
return counts
def _count_arr(self, arr: np.ndarray, length: int) -> np.ndarray:
counts = np.zeros(length, dtype=int)
for idx in arr:
counts[idx] += 1
return counts
def _lower_confidence_bound(self, NA: int, N: int, alpha: float) -> float:
""" Returns a (1 - alpha) lower confidence bound on a bernoulli proportion.
This function uses the Clopper-Pearson method.
:param NA: the number of "successes"
:param N: the number of total draws
:param alpha: the confidence level
:return: a lower bound on the binomial proportion which holds true w.p at least (1 - alpha) over the samples
"""
return proportion_confint(NA, N, alpha=2 * alpha, method="beta")[0]