def project(self, x: ep.Tensor, x0: ep.Tensor,
epsilon: float) -> ep.Tensor:
# based on https://github.com/ftramer/MultiRobustness/blob/ad41b63235d13b1b2a177c5f270ab9afa74eee69/pgd_attack.py#L110
delta = flatten(x - x0)
norms = delta.norms.l1(axis=-1)
if (norms <= epsilon).all():
return x
n, d = delta.shape
abs_delta = abs(delta)
mu = -ep.sort(-abs_delta, axis=-1)
cumsums = mu.cumsum(axis=-1)
js = 1.0 / ep.arange(x, 1, d + 1).astype(x.dtype)
temp = mu - js * (cumsums - epsilon)
guarantee_first = ep.arange(x, d).astype(x.dtype) / d
# guarantee_first are small values (< 1) that we add to the boolean
# tensor (only 0 and 1) to break the ties and always return the first
# argmin, i.e. the first value where the boolean tensor is 0
# (otherwise, this is not guaranteed on GPUs, see e.g. PyTorch)
rho = ep.argmin((temp > 0).astype(x.dtype) + guarantee_first, axis=-1)
theta = 1.0 / (1 + rho.astype(x.dtype)) * (cumsums[range(n), rho] -
epsilon)
delta = delta.sign() * ep.maximum(abs_delta - theta[..., ep.newaxis],
0)
delta = delta.reshape(x.shape)
return x0 + delta
def project_onto_l1_ball(x: ep.Tensor, eps: ep.Tensor) -> ep.Tensor:
"""Computes Euclidean projection onto the L1 ball for a batch. [#Duchi08]_
Adapted from the pytorch version by Tony Duan:
https://gist.github.com/tonyduan/1329998205d88c566588e57e3e2c0c55
Args:
x: Batch of arbitrary-size tensors to project, possibly on GPU
eps: radius of l-1 ball to project onto
References:
..[#Duchi08] Efficient Projections onto the l1-Ball for Learning in High Dimensions
John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
International Conference on Machine Learning (ICML 2008)
"""
original_shape = x.shape
x = flatten(x)
mask = (ep.norms.l1(x, axis=1) <= eps).astype(x.dtype).expand_dims(1)
mu = ep.flip(ep.sort(ep.abs(x)), axis=-1).astype(x.dtype)
cumsum = ep.cumsum(mu, axis=-1)
arange = ep.arange(x, 1, x.shape[1] + 1).astype(x.dtype)
rho = (ep.max(
((mu * arange >
(cumsum - eps.expand_dims(1)))).astype(x.dtype) * arange,
axis=-1,
) - 1)
# samples already under norm will have to select
rho = ep.maximum(rho, 0)
theta = (cumsum[ep.arange(x, x.shape[0]),
rho.astype(ep.arange(x, 1).dtype)] - eps) / (rho + 1.0)
proj = (ep.abs(x) - theta.expand_dims(1)).clip(min_=0, max_=ep.inf)
x = mask * x + (1 - mask) * proj * ep.sign(x)
return x.reshape(original_shape)
def project_onto_l1_ball(x: ep.Tensor, eps: ep.Tensor):
"""
Compute Euclidean projection onto the L1 ball for a batch.
min ||x - u||_2 s.t. ||u||_1 <= eps
Inspired by the corresponding numpy version by Adrien Gaidon.
Adapted from the pytorch version by Tony Duan: https://gist.github.com/tonyduan/1329998205d88c566588e57e3e2c0c55
Parameters
----------
x: (batch_size, *) torch array
batch of arbitrary-size tensors to project, possibly on GPU
eps: float
radius of l-1 ball to project onto
Returns
-------
u: (batch_size, *) torch array
batch of projected tensors, reshaped to match the original
Notes
-----
The complexity of this algorithm is in O(dlogd) as it involves sorting x.
References
----------
[1] Efficient Projections onto the l1-Ball for Learning in High Dimensions
John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
International Conference on Machine Learning (ICML 2008)
"""
original_shape = x.shape
x = flatten(x)
mask = (ep.norms.l1(x, axis=1) < eps).astype(x.dtype).expand_dims(1)
mu = ep.flip(ep.sort(ep.abs(x)), axis=-1)
cumsum = ep.cumsum(mu, axis=-1)
arange = ep.arange(x, 1, x.shape[1] + 1)
rho = ep.max(
(mu * arange > (cumsum - eps.expand_dims(1))) * arange, axis=-1) - 1
theta = (cumsum[ep.arange(x, x.shape[0]), rho] - eps) / (rho + 1.0)
proj = (ep.abs(x) - theta.expand_dims(1)).clip(min_=0, max_=ep.inf)
x = mask * x + (1 - mask) * proj * ep.sign(x)
return x.reshape(original_shape)
def test_sort(dummy: Tensor) -> Tensor:
t = -ep.arange(dummy, 6).float32().reshape((2, 3))
return ep.sort(t)