-
Notifications
You must be signed in to change notification settings - Fork 3
/
core.py
112 lines (96 loc) · 4.91 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
import torch
from scipy.stats import norm, binom_test
import numpy as np
from math import ceil
from statsmodels.stats.proportion import proportion_confint
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]