Spaces:
Runtime error
Runtime error
File size: 3,670 Bytes
fc0ff8f |
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 |
# Code taken from https://github.com/chs20/RobustVLM/tree/main
import torch
from vlm_eval.attacks.utils import project_perturbation, normalize_grad
class PGD:
"""
Minimize or maximize given loss
"""
def __init__(self, forward, norm, eps, mode='min', mask_out='context', image_space=True):
self.model = forward
self.norm = norm
self.eps = eps
self.momentum = 0.9
self.mode = mode
self.mask_out = mask_out
self.image_space = image_space
def perturb(self, data_clean, iterations, stepsize, perturbation=None, verbose=False, return_loss=False):
if self.image_space:
# make sure data is in image space
assert torch.max(data_clean) < 1. + 1e-6 and torch.min(data_clean) > -1e-6 # todo
if perturbation is None:
perturbation = torch.zeros_like(data_clean, requires_grad=True)
mask = self._set_mask(data_clean)
velocity = torch.zeros_like(data_clean)
for i in range(iterations):
perturbation.requires_grad_()
with torch.enable_grad():
loss = self.model(data_clean + perturbation)
# print 10 times in total and last iteration
if verbose and (i % (iterations // 10 + 1) == 0 or i == iterations - 1):
print(f'[iteration] {i} [loss] {loss.item()}')
with torch.no_grad():
gradient = torch.autograd.grad(loss, perturbation)[0]
gradient = mask * gradient
if gradient.isnan().any(): #
print(f'attention: nan in gradient ({gradient.isnan().sum()})') #
gradient[gradient.isnan()] = 0.
# normalize
gradient = normalize_grad(gradient, p=self.norm)
# momentum
velocity = self.momentum * velocity + gradient
velocity = normalize_grad(velocity, p=self.norm)
# update
if self.mode == 'min':
perturbation = perturbation - stepsize * velocity
elif self.mode == 'max':
perturbation = perturbation + stepsize * velocity
else:
raise ValueError(f'Unknown mode: {self.mode}')
# project
perturbation = project_perturbation(perturbation, self.eps, self.norm)
if self.image_space:
perturbation = torch.clamp(
data_clean + perturbation, 0, 1
) - data_clean # clamp to image space
assert torch.max(data_clean + perturbation) < 1. + 1e-6 and torch.min(
data_clean + perturbation
) > -1e-6
assert not perturbation.isnan().any()
# assert (ctorch.compute_norm(perturbation, p=self.norm) <= self.eps + 1e-6).all()
# todo return best perturbation
# problem is that model currently does not output expanded loss
if return_loss:
return data_clean + perturbation.detach(), loss
else:
return data_clean + perturbation.detach()
def _set_mask(self, data):
mask = torch.ones_like(data)
if self.mask_out == 'context':
mask[:, :-1, ...] = 0
elif self.mask_out == 'query':
mask[:, -1, ...] = 0
elif isinstance(self.mask_out, int):
mask[:, self.mask_out, ...] = 0
elif self.mask_out is None:
pass
else:
raise NotImplementedError(f'Unknown mask_out: {self.mask_out}')
return mask
|