import torch from torch import nn from .. import losses import ignite.distributed as idist import torch_optimizer from tqdm import tqdm import matplotlib.pyplot as plt from torch.nn import functional as F import os import shutil from modelguidedattacks.cls_models.registry import MMPretrainVisualTransformerWrapper from modelguidedattacks.data.imagenet_metadata import imgnet_idx_to_name class Unguided(nn.Module): def __init__(self, model: nn.Module, config, optimizer=torch.optim.AdamW, seed=0, iterations=1000, loss_fn=losses.CVXProjLoss, lr=1e-3, binary_search_steps=1, topk_loss_coef_upper=10., topk_loss_coef_lower=0.) -> None: super().__init__() self.guided = False self.model = model self.seed = seed self.iterations = iterations self.loss = loss_fn() self.optimizer = optimizer self.lr = lr self.binary_search_steps = binary_search_steps self.topk_loss_coef_upper = topk_loss_coef_upper self.topk_loss_coef_lower = topk_loss_coef_lower self.config = config def surject_perturbation(self, x, max_norm=5.): x_shape = x.shape x = x.flatten(1) x_norm = x.norm(dim=-1) x_unit = x / x_norm[:, None] x_norm_outside = x_norm > max_norm x_norm_outside = x_norm_outside.expand_as(x) x = torch.where(x_norm_outside, x_unit*max_norm, x) return x.view(x_shape) @torch.enable_grad() def attack(self, x, attack_targets, gt_labels, topk_coefs): """ For a given set of topk coefficients, this function computes best energy attack in the given number of iterations and configuration x: [B, C, H, W] [0-1 for colors] attack_targets: [B, K] (long) gt_labels: [B] (long) topk_coefs: [B] (floats) """ topk_coefs = topk_coefs.clone() K = attack_targets.shape[-1] x_perturbation = nn.Parameter(torch.randn(x.shape, device=x.device)*2e-3) optimizer = self.optimizer([x_perturbation], lr=self.lr) precomputed_state = self.loss.precompute(attack_targets, gt_labels, self.config) with torch.no_grad(): prediction_logits_0, prediction_feats_0 \ = self.model(x, return_features=True) best_perturbations = torch.zeros_like(x) # [B, 3, H, W] has_successful_attack = torch.zeros(x.shape[0], dtype=torch.long, device=x.device) # [B] best_energy = torch.full((x.shape[0],), float('inf'), device=x.device) # [B] pbar = tqdm(range(self.iterations)) for i in pbar: if i == self.config.opt_warmup_its: # Reset optimizer state optimizer = self.optimizer([x_perturbation], lr=self.lr) x_perturbed = x + x_perturbation#self.surject_perturbation(x_perturbation) prediction_logits, prediction_feats = self.model(x_perturbed, return_features=True) pred_classes = prediction_logits.argsort(dim=-1, descending=True) # [B, C] attack_successful = (pred_classes[:, :K] == attack_targets).all(dim=-1) # [B] attack_energy = x_perturbation.flatten(1).norm(dim=-1) # [B] attack_improved = attack_successful & (attack_energy <= best_energy) best_perturbations[attack_improved] = x_perturbation[attack_improved] has_successful_attack[attack_improved] = True best_energy[attack_improved] = attack_energy[attack_improved] loss = self.loss(logits_pred=prediction_logits, feats_pred=prediction_feats, feats_pred_0=prediction_feats_0, attack_targets=attack_targets, model=self.model, **precomputed_state) loss = loss * topk_coefs loss = loss.sum() pbar.set_description(f"Loss: {loss.item():.3f}") loss = loss + x_perturbation.flatten(1).square().sum() optimizer.zero_grad() loss.backward() optimizer.step() # If we were successfull let's start taking the norm down topk_coefs[attack_improved] *= 0.75 # Project perturbation to be within image limits with torch.no_grad(): x_perturbed = x + x_perturbation x_perturbed = x_perturbed.clamp_(min=0., max=1.) x_perturbation.data = x_perturbed - x x_perturbed_best = x + best_perturbations prediction_logits, prediction_feats = self.model(x_perturbed_best, return_features=True) if self.config.dump_plots: if os.path.isdir(self.config.plot_out): shutil.rmtree(self.config.plot_out) if has_successful_attack.any(): def dump_random_map(): os.makedirs(self.config.plot_out, exist_ok=True) # selected_idx = best_energy.argmin() successful_idxs = has_successful_attack.nonzero()[:, 0] if self.config.plot_idx == "find": selected_idx = successful_idxs[torch.randperm(len(successful_idxs))[0]] # selected_idx = best_energy.argmin() else: selected_idx = int(self.config.plot_idx) print ("Selected idx", selected_idx) top_classes = prediction_logits_0[selected_idx].argsort(dim=-1, descending=True) attack_targets_selected = attack_targets[selected_idx] def imgnet_names(idxs): return [imgnet_idx_to_name[int(idx)].split(",")[0] for idx in idxs] top_class_names = imgnet_names(top_classes)[:K] attack_targets_selected_names = imgnet_names(attack_targets_selected) def plot_attn_map(attn_map): attn_map = attn_map[0].mean(dim=0)[1:] # [196] get class tokens attn_map = attn_map.view(14, 14) attn_map = F.interpolate( attn_map[None, None], x.shape[-2:], mode="bilinear" ).view(x.shape[-2:]) plt.imshow(attn_map.detach().cpu(), alpha=0.5) plt.figure() plt.imshow(x[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu()) plt.axis("off") plt.savefig(f"{self.config.plot_out}/clean_image.png", bbox_inches="tight", pad_inches=0) plt.figure() plt.imshow(x_perturbed_best[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu()) plt.axis("off") plt.savefig(f"{self.config.plot_out}/perturbed_image.png", bbox_inches="tight", pad_inches=0) plt.figure() plt.imshow(best_perturbations[selected_idx].mean(dim=0).abs().detach().cpu(), cmap="hot") plt.colorbar() plt.savefig(f"{self.config.plot_out}/perturbation.png", bbox_inches="tight") if isinstance(self.model, MMPretrainVisualTransformerWrapper): attn_maps_clean = self.model.get_attention_maps(x)[-1][selected_idx] attn_maps_attacked = self.model.get_attention_maps(x_perturbed_best)[-1][selected_idx] plt.figure() plt.imshow(x[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu()) plot_attn_map(attn_maps_clean) plt.axis("off") plt.savefig(f"{self.config.plot_out}/clean_map.png", bbox_inches="tight", pad_inches=0) plt.figure() plt.imshow(x[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu()) plot_attn_map(attn_maps_attacked) plt.axis("off") plt.savefig(f"{self.config.plot_out}/attacked_map.png", bbox_inches="tight", pad_inches=0) with open(f'{self.config.plot_out}/clean_classes_names.txt', 'w') as f: f.write(", ".join(top_class_names)) with open(f'{self.config.plot_out}/attack_targets_names.txt', 'w') as f: f.write(", ".join(attack_targets_selected_names)) with open(f'{self.config.plot_out}/clean_classes_names.txt', 'w') as f: f.write(", ".join(top_class_names)) with open(f'{self.config.plot_out}/selected_idx.txt', 'w') as f: if isinstance(selected_idx, torch.Tensor): selected_idx = selected_idx.item() f.write(str(selected_idx)) with open(f'{self.config.plot_out}/energy.txt', 'w') as f: f.write(str(best_energy[selected_idx].item())) C = prediction_logits_0.shape[-1] class_idxs = torch.arange(C) + 1 clean_probs = prediction_logits_0[selected_idx].detach().cpu().softmax(dim=-1) attacked_probs = prediction_logits[selected_idx].detach().cpu().softmax(dim=-1) def label_classes(bars): adjusted_heights = {} for i, cls_idx in enumerate(attack_targets_selected.tolist()): bar = bars[cls_idx] height = bar.get_height() ann_x = bar.get_x() + bar.get_width() rotation = 90 font_size = 10 max_neighboring_height = -1 for other_cls_idx in attack_targets_selected.tolist(): if abs(cls_idx - other_cls_idx) <= 40 and cls_idx != other_cls_idx: if other_cls_idx in adjusted_heights and adjusted_heights[other_cls_idx] > max_neighboring_height: max_neighboring_height = adjusted_heights[other_cls_idx] if max_neighboring_height > 0: height = max_neighboring_height + 0.05 adjusted_heights[cls_idx] = height plt.text(ann_x, height, f"[{i}]", rotation=rotation, ha='center', va='bottom', fontsize=font_size, color='red')#.get_bbox_patch().get_height() plt.figure() bars_clean = plt.bar(class_idxs, clean_probs, width=4) plt.ylim(0,1) label_classes(bars_clean) plt.savefig(f"{self.config.plot_out}/clean_probs.png", bbox_inches="tight", pad_inches=0) plt.figure() bars_attacked = plt.bar(class_idxs, attacked_probs, width=4) plt.ylim(0,1) label_classes(bars_attacked) plt.savefig(f"{self.config.plot_out}/attacked_probs.png", bbox_inches="tight", pad_inches=0) print ("Idx", selected_idx) print (best_energy[selected_idx]) print ("Finished plotting") dump_random_map() import sys sys.exit(1) print ("Dumped attention map") return prediction_logits, best_perturbations, best_energy def forward(self, x, attack_targets, gt_labels): """ This function is in charge of performing a binary search through topk loss coefficients and running attacks on each. """ B = x.shape[0] device = x.device topk_coefs_lower = torch.full((B,), fill_value=self.topk_loss_coef_lower, device=device, dtype=torch.float) topk_coefs_upper = torch.full((B,), fill_value=self.topk_loss_coef_upper, device=device, dtype=torch.float) best_perturbations = torch.zeros_like(x) # [B, 3, H, W] best_energy = torch.full((B,), float('inf'), device=device) # [B] best_prediction_logits = None for search_step_i in range(self.binary_search_steps): if x.device.index is None or x.device.index == 0: print ("Running binary search step", search_step_i + 1) current_topk_coefs = (topk_coefs_lower + topk_coefs_upper) / 2 current_logits, current_perturbations, current_energy = \ self.attack(x, attack_targets, gt_labels, current_topk_coefs) current_attack_suceeded = ~torch.isinf(current_energy) update_mask = current_energy < best_energy best_perturbations[update_mask] = current_perturbations[update_mask] best_energy[update_mask] = current_energy[update_mask] if best_prediction_logits is None: best_prediction_logits = current_logits.clone() else: best_prediction_logits[update_mask] = current_logits[update_mask] # If we fail to attack, we must increase our topk coef topk_coefs_lower[~current_attack_suceeded] = current_topk_coefs[~current_attack_suceeded] # If we succeed, we must lower to seek a more frugal attack topk_coefs_upper[current_attack_suceeded] = current_topk_coefs[current_attack_suceeded] idist.barrier() return best_prediction_logits, best_perturbations