| 234 | |
| 235 | |
| 236 | class AdversarialInputAttacker: |
| 237 | def __init__(self, model: List[torch.nn.Module], epsilon=16 / 255, norm="Linf"): |
| 238 | assert norm in ["Linf", "L2"] |
| 239 | self.norm = norm |
| 240 | self.epsilon = epsilon |
| 241 | self.models = model |
| 242 | self.init() |
| 243 | self.model_distribute() |
| 244 | self.device = torch.device("cuda") |
| 245 | self.n = len(self.models) |
| 246 | |
| 247 | @abstractmethod |
| 248 | def attack(self, *args, **kwargs): |
| 249 | pass |
| 250 | |
| 251 | def __call__(self, *args, **kwargs): |
| 252 | return self.attack(*args, **kwargs) |
| 253 | |
| 254 | def model_distribute(self): |
| 255 | """ |
| 256 | make each model on one gpu |
| 257 | :return: |
| 258 | """ |
| 259 | num_gpus = torch.cuda.device_count() |
| 260 | models_each_gpu = ceil(len(self.models) / num_gpus) |
| 261 | for i, model in enumerate(self.models): |
| 262 | model.to(torch.device(f"cuda:{num_gpus - 1 - i // models_each_gpu}")) |
| 263 | model.device = torch.device(f"cuda:{num_gpus - 1 - i // models_each_gpu}") |
| 264 | |
| 265 | def init(self): |
| 266 | # set the model parameters requires_grad is False |
| 267 | for model in self.models: |
| 268 | model.requires_grad_(False) |
| 269 | model.eval() |
| 270 | |
| 271 | def to(self, device: torch.device): |
| 272 | for model in self.models: |
| 273 | model.to(device) |
| 274 | model.device = device |
| 275 | self.device = device |
| 276 | |
| 277 | def clamp(self, x: Tensor, ori_x: Tensor) -> Tensor: |
| 278 | B = x.shape[0] |
| 279 | if self.norm == "Linf": |
| 280 | x = torch.clamp(x, min=ori_x - self.epsilon, max=ori_x + self.epsilon) |
| 281 | elif self.norm == "L2": |
| 282 | difference = x - ori_x |
| 283 | distance = torch.norm(difference.view(B, -1), p=2, dim=1) |
| 284 | mask = distance > self.epsilon |
| 285 | if torch.sum(mask) > 0: |
| 286 | difference[mask] = difference[mask] / distance[mask].view(torch.sum(mask), 1, 1, 1) * self.epsilon |
| 287 | x = ori_x + difference |
| 288 | x = torch.clamp(x, min=0, max=1) |
| 289 | return x |
| 290 | |
| 291 | |
| 292 | class SpectrumSimulationAttack(AdversarialInputAttacker): |
nothing calls this directly
no outgoing calls
no test coverage detected