-
Notifications
You must be signed in to change notification settings - Fork 2
/
core.py
163 lines (139 loc) · 7.8 KB
/
core.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
from math import ceil
import numpy as np
from scipy.stats import norm, binom_test
from statsmodels.stats.proportion import proportion_confint
import torch
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, sample_id:int, alpha: float, batch_size: int, clustering_method='none') -> (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, n0_predictions = self._sample_noise(x, n0, batch_size, clustering_method, sample_id)
# 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, n_predictions = self._sample_noise(x, n, batch_size, clustering_method, sample_id)
# 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, n0_predictions, n_predictions
else:
radius = self.sigma * norm.ppf(pABar)
return cAHat, radius, n0_predictions, n_predictions
def certify_noapproximate(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
top2 = counts_estimation.argsort()[::-1][:2]
nA = counts_estimation[top2[0]].item()
nB = counts_estimation[top2[1]].item()
pABar = self._lower_confidence_bound(nA, n, alpha)
pBBar = self._upper_confidence_bound(nB, n, alpha)
if pABar < 0.5:
return Smooth.ABSTAIN, 0.0
else:
radius = self.sigma/2 * (norm.ppf(pABar) - norm.ppf(pBBar))
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, clustering_method='none', sample_id=None) -> 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():
predictions_all = np.array([], dtype=int)
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
if clustering_method == 'classifier':
predictions = self.base_classifier(batch + noise, sample_id).argmax(1)
predictions = predictions.view(this_batch_size,-1).cpu().numpy()
count_max_list = np.zeros(this_batch_size,dtype=int)
for i in range(this_batch_size):
count_max = max(list(predictions[i]),key=list(predictions[i]).count)
count_max_list[i] = count_max
counts += self._count_arr(count_max_list, self.num_classes)
else:
predictions = self.base_classifier(batch + noise, sample_id).argmax(1)
counts += self._count_arr(predictions.cpu().numpy(), self.num_classes)
predictions_all = np.hstack((predictions_all, predictions.cpu().numpy()))
return counts, predictions_all
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]
def _upper_confidence_bound(self, NA: int, N: int, alpha: float) -> float:
return proportion_confint(NA, N, alpha=2 * alpha, method="beta")[1]