"""
Randomized smooothing certification. Based on code from
https://github.com/locuslab/smoothing
"""
import torch
import torch.nn
import torch.nn.functional as F
import numpy as np
from scipy.stats import norm, binom_test
from statsmodels.stats.proportion import proportion_confint
from math import ceil
def quick_smoothing(model, x, y, sigma=1.0, eps=1.0,
num_smooth=100, batch_size=1000,
softmax_temperature=100.0,
detailed_output=False):
"""Quick and dirty randomized smoothing 'certification', without proper
confidence bounds. We use it only to monitor training."""
x_noise = x.view(1, *x.shape) + sigma * torch.randn(
num_smooth, *x.shape).cuda()
x_noise = x_noise.view(-1, *x.shape[1:])
# by setting a high softmax temperature, we are effectively using the
# randomized smoothing approach as originally defined
# it will be interesting to see if lower temperatures help
preds = torch.cat([F.softmax(softmax_temperature * model(batch), dim=-1)
for batch in torch.split(x_noise, batch_size)])
preds = preds.view(num_smooth, x.shape[0], -1).mean(dim=0)
p_max, y_pred = preds.max(dim=-1)
correct = (y_pred == y).cpu().numpy().astype('int64')
radii = (sigma + 1e-16) * norm.ppf(p_max.cpu().numpy())
err = (1 - correct).sum()
robust_err = (1 - correct * (radii >= eps)).sum()
if not detailed_output:
return err, robust_err
else:
return correct, radii
def add_noise(x: torch.Tensor, noise_sd: float) -> torch.Tensor:
noise = torch.randn_like(x)
# if args.distribution != 'gaussian':
# noise_flat = noise.view(noise.shape[0], -1)
# noise_flat /= noise_flat.norm(dim=-1).view(-1, 1)
return x + noise * noise_sd
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, pABar, 0.0, counts_estimation
else:
radius = self.sigma * norm.ppf(pABar)
return cAHat, pABar, radius, counts_estimation
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))
batch_noise = add_noise(batch, 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]
