1{
2  "op": "replace",
3  "filename": "torchattacks/bench/custom_attack.py",
4  "content": "import torch.nn.functional as F\n\n# =====================================================================\n# EDITABLE: implement run_attack below\n# =====================================================================\ndef run_attack(\n    model: nn.Module,\n    images: torch.Tensor,\n    labels: torch.Tensor,\n    pixels: int,\n    device: torch.device,\n    n_classes: int,\n) -> torch.Tensor:\n    \"\"\"\n    Sparse L0 adversarial attack based on multi-target support refinement.\n\n    The attack alternates between:\n      1) selecting a k-pixel support with a first-order boundary-gain score,\n      2) pushing the selected pixels toward loss-increasing box vertices,\n      3) refining values on the fixed support with masked gradient ascent.\n\n    We restart this dynamic support search on several strong target classes and\n    keep the adversarial image with the largest untargeted decision margin.\n    \"\"\"\n    model.eval()\n    images = images.to(device)\n    labels = labels.to(device)\n    adv_images = images.clone()\n    batch_size = min(128, images.size(0))\n\n    def decision_margin(logits: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n        ar = torch.arange(logits.size(0), device=logits.device)\n        true_logit = logits[ar, y]\n        masked = logits.clone()\n        masked[ar, y] = -1e9\n        other_logit = masked.max(dim=1).values\n        return other_logit - true_logit\n\n    def objective(logits: torch.Tensor, y: torch.Tensor, target: torch.Tensor | None) -> torch.Tensor:\n        ar = torch.arange(logits.size(0), device=logits.device)\n        true_logit = logits[ar, y]\n        masked = logits.clone()\n        masked[ar, y] = -1e9\n        other_logit = masked.max(dim=1).values\n        if target is None:\n            return torch.logsumexp(masked, dim=1) - true_logit\n        logp = F.log_softmax(logits, dim=1)\n        return 0.8 * (logp[ar, target] - logp[ar, y]) + 0.2 * (other_logit - true_logit)\n\n    def attack_once(x: torch.Tensor, y: torch.Tensor, target: torch.Tensor | None) -> torch.Tensor:\n        b, _, h, w = x.shape\n        adv = x.clone()\n        mask = torch.zeros(b, 1, h, w, dtype=torch.bool, device=x.device)\n        momentum = torch.zeros_like(x)\n        k = min(pixels, h * w)\n\n        for _ in range(4):\n            adv.requires_grad_(True)\n            logits = model(adv)\n            loss = objective(logits, y, target).sum()\n            grad = torch.autograd.grad(loss, adv)[0]\n            grad = grad / (grad.abs().mean(dim=(1, 2, 3), keepdim=True) + 1e-12)\n            momentum = 0.75 * momentum + grad\n\n            gain = momentum.clamp(min=0) * (1.0 - adv.detach()) + (-momentum).clamp(min=0) * adv.detach()\n            gain = gain.sum(dim=1)\n            gain = gain + 0.15 * gain.mean(dim=(1, 2), keepdim=True) * mask.squeeze(1).float()\n            flat_gain = gain.flatten(1)\n            topk = flat_gain.topk(k, dim=1).indices\n            mask = torch.zeros_like(flat_gain, dtype=torch.bool).scatter(1, topk, True).view(b, 1, h, w)\n\n            extreme = torch.where(momentum >= 0, torch.ones_like(adv), torch.zeros_like(adv))\n            adv = adv.detach()\n            adv = torch.where(mask, adv + 0.7 * (extreme - adv), x).clamp(0.0, 1.0)\n\n            adv.requires_grad_(True)\n            logits = model(adv)\n            loss = objective(logits, y, target).sum()\n            grad = torch.autograd.grad(loss, adv)[0]\n            adv = (adv.detach() + 0.35 * grad.sign() * mask).clamp(0.0, 1.0)\n            adv = torch.where(mask, adv, x)\n\n        for _ in range(2):\n            adv.requires_grad_(True)\n            logits = model(adv)\n            loss = objective(logits, y, target).sum()\n            grad = torch.autograd.grad(loss, adv)[0]\n            adv = (adv.detach() + 0.2 * grad.sign() * mask).clamp(0.0, 1.0)\n            adv = torch.where(mask, adv, x)\n        return adv.detach()\n\n    with torch.enable_grad():\n        for start in range(0, images.size(0), batch_size):\n            end = start + batch_size\n            x = images[start:end]\n            y = labels[start:end]\n\n            with torch.no_grad():\n                clean_logits = model(x)\n                order = clean_logits.argsort(dim=1, descending=True)\n                alt_targets = order[order.ne(y[:, None])].view(x.size(0), -1)\n                n_targets = min(3 if n_classes <= 10 else 4, alt_targets.size(1))\n                best_margin = decision_margin(clean_logits, y)\n\n            best_adv = x.clone()\n            success = best_margin > 0\n\n            for t in range(n_targets + 1):\n                active_idx = (~success).nonzero(as_tuple=False).squeeze(1)\n                if active_idx.numel() == 0:\n                    break\n                xa, ya = x[active_idx], y[active_idx]\n                target = alt_targets[active_idx, t] if t < n_targets else None\n                cand_adv = attack_once(xa, ya, target)\n\n                with torch.no_grad():\n                    cand_logits = model(cand_adv)\n                    cand_margin = decision_margin(cand_logits, ya)\n\n                improve = cand_margin > best_margin[active_idx]\n                if improve.any():\n                    chosen = active_idx[improve]\n                    best_adv[chosen] = cand_adv[improve]\n                    best_margin[chosen] = cand_margin[improve]\n                success[active_idx] |= cand_margin > 0\n\n            adv_images[start:end] = best_adv\n\n    return adv_images.clamp(0.0, 1.0)\n\n# =====================================================================\n# END EDITABLE REGION\n# =====================================================================",
5  "after_line": 0,
6  "start_line": 3,
7  "end_line": 26
8}

1{
2  "op": "replace",
3  "filename": "torchattacks/bench/custom_attack.py",
4  "content": "import torch.nn.functional as F\n\n# =====================================================================\n# EDITABLE: implement run_attack below\n# =====================================================================\ndef run_attack(\n    model: nn.Module,\n    images: torch.Tensor,\n    labels: torch.Tensor,\n    pixels: int,\n    device: torch.device,\n    n_classes: int,\n) -> torch.Tensor:\n    \"\"\"\n    Expand-prune sparse attack with support replacement.\n\n    Core idea:\n      1) search over an expanded support (> budget) using box-aware gradient gain,\n      2) prune back to the exact pixel budget using contribution scores,\n      3) refine the support with one-pixel replacement moves,\n      4) repeat across several high-probability target classes and keep the best.\n    \"\"\"\n    model.eval()\n    images = images.to(device)\n    labels = labels.to(device)\n    adv_images = images.clone()\n    batch_size = min(128, images.size(0))\n\n    def decision_margin(logits: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n        row = torch.arange(logits.size(0), device=logits.device)\n        true_logit = logits[row, y]\n        masked = logits.clone()\n        masked[row, y] = -1e9\n        other_logit = masked.max(dim=1).values\n        return other_logit - true_logit\n\n    def objective(\n        logits: torch.Tensor,\n        y: torch.Tensor,\n        target: torch.Tensor | None,\n    ) -> torch.Tensor:\n        row = torch.arange(logits.size(0), device=logits.device)\n        true_logit = logits[row, y]\n        masked = logits.clone()\n        masked[row, y] = -1e9\n        other_logit = masked.max(dim=1).values\n        if target is None:\n            return torch.logsumexp(masked, dim=1) - true_logit\n        logp = F.log_softmax(logits, dim=1)\n        tgt_w = 0.9 if n_classes > 10 else 0.8\n        return tgt_w * (logp[row, target] - logp[row, y]) + (1.0 - tgt_w) * (other_logit - true_logit)\n\n    def boundary_gain(grad: torch.Tensor, point: torch.Tensor) -> torch.Tensor:\n        return (grad.clamp(min=0) * (1.0 - point) + (-grad).clamp(min=0) * point).sum(dim=1)\n\n    def choose_best(\n        candidates: list[torch.Tensor],\n        y: torch.Tensor,\n        target: torch.Tensor | None,\n    ) -> torch.Tensor:\n        if len(candidates) == 1:\n            return candidates[0]\n        bsz = candidates[0].size(0)\n        stacked = torch.cat(candidates, dim=0)\n        logits = model(stacked)\n        y_rep = y.repeat(len(candidates))\n        tgt_rep = None if target is None else target.repeat(len(candidates))\n        scores = objective(logits, y_rep, tgt_rep).view(len(candidates), bsz).t()\n        best_idx = scores.argmax(dim=1)\n        best = candidates[0].clone()\n        for i, cand in enumerate(candidates):\n            pick = best_idx == i\n            if pick.any():\n                best[pick] = cand[pick]\n        return best\n\n    def masked_step(\n        x: torch.Tensor,\n        adv: torch.Tensor,\n        y: torch.Tensor,\n        target: torch.Tensor | None,\n        mask: torch.Tensor,\n        step_size: float,\n    ) -> tuple[torch.Tensor, torch.Tensor]:\n        adv.requires_grad_(True)\n        logits = model(adv)\n        loss = objective(logits, y, target).sum()\n        grad = torch.autograd.grad(loss, adv)[0]\n        grad = grad / (grad.abs().mean(dim=(1, 2, 3), keepdim=True) + 1e-12)\n        proposal = (adv.detach() + step_size * grad.sign() * mask).clamp(0.0, 1.0)\n        proposal = torch.where(mask, proposal, x)\n        extreme = torch.where(grad >= 0, torch.ones_like(adv), torch.zeros_like(adv))\n        corner = torch.where(mask, adv.detach() + 0.5 * (extreme - adv.detach()), x).clamp(0.0, 1.0)\n        best = choose_best([adv.detach(), proposal, corner], y, target)\n        return torch.where(mask, best, x).detach(), grad.detach()\n\n    def attack_once(x: torch.Tensor, y: torch.Tensor, target: torch.Tensor | None) -> torch.Tensor:\n        bsz, _, h, w = x.shape\n        k = min(pixels, h * w)\n        k_search = min(h * w, max(k, 2 * k))\n        adv = x.clone()\n        mask = torch.zeros(bsz, 1, h, w, dtype=torch.bool, device=x.device)\n        momentum = torch.zeros_like(x)\n\n        for _ in range(3):\n            adv.requires_grad_(True)\n            logits = model(adv)\n            loss = objective(logits, y, target).sum()\n            grad = torch.autograd.grad(loss, adv)[0]\n            grad = grad / (grad.abs().mean(dim=(1, 2, 3), keepdim=True) + 1e-12)\n            momentum = 0.8 * momentum + grad\n\n            gain = boundary_gain(momentum, adv.detach()).flatten(1)\n            prev = mask.flatten(1).float()\n            gain = gain + 0.1 * gain.mean(dim=1, keepdim=True) * prev\n            top = gain.topk(k_search, dim=1).indices\n            flat_mask = torch.zeros_like(gain, dtype=torch.bool)\n            flat_mask.scatter_(1, top, True)\n            mask = flat_mask.view(bsz, 1, h, w)\n\n            extreme = torch.where(momentum >= 0, torch.ones_like(adv), torch.zeros_like(adv))\n            c1 = torch.where(mask, adv.detach() + 0.55 * (extreme - adv.detach()), x).clamp(0.0, 1.0)\n            c2 = torch.where(mask, adv.detach() + 0.85 * (extreme - adv.detach()), x).clamp(0.0, 1.0)\n            c3 = torch.where(mask, extreme, x)\n            adv = choose_best([adv.detach(), c1, c2, c3], y, target)\n            adv = torch.where(mask, adv, x)\n            adv, step_grad = masked_step(x, adv, y, target, mask, 0.28)\n            momentum = 0.7 * momentum + 0.3 * step_grad\n\n        adv.requires_grad_(True)\n        logits = model(adv)\n        loss = objective(logits, y, target).sum()\n        grad = torch.autograd.grad(loss, adv)[0]\n        grad = grad / (grad.abs().mean(dim=(1, 2, 3), keepdim=True) + 1e-12)\n\n        if k_search > k:\n            contrib = ((adv.detach() - x).abs() * (0.25 + grad.abs())).sum(dim=1).flatten(1)\n            flat_mask = mask.flatten(1)\n            contrib = contrib.masked_fill(~flat_mask, -1e9)\n            keep = contrib.topk(k, dim=1).indices\n            pruned = torch.zeros_like(flat_mask)\n            pruned.scatter_(1, keep, True)\n            mask = pruned.view(bsz, 1, h, w)\n            adv = torch.where(mask, adv.detach(), x)\n            adv, grad = masked_step(x, adv, y, target, mask, 0.24)\n\n        for _ in range(2):\n            flat_mask = mask.flatten(1)\n            outside = boundary_gain(grad, adv.detach()).flatten(1).masked_fill(flat_mask, -1e9)\n            inside = ((adv.detach() - x).abs() * (0.25 + grad.abs())).sum(dim=1).flatten(1)\n            inside = inside.masked_fill(~flat_mask, 1e9)\n            best_out = outside.argmax(dim=1, keepdim=True)\n            worst_in = inside.argmin(dim=1, keepdim=True)\n            row = torch.arange(bsz, device=x.device)[:, None]\n            swapped = flat_mask.clone()\n            swapped[row, worst_in] = False\n            swapped[row, best_out] = True\n            mask = swapped.view(bsz, 1, h, w)\n            adv = torch.where(mask, adv.detach(), x)\n\n            extreme = torch.where(grad >= 0, torch.ones_like(adv), torch.zeros_like(adv))\n            moved = torch.where(mask, adv + 0.65 * (extreme - adv), x).clamp(0.0, 1.0)\n            adv = choose_best([adv.detach(), moved], y, target)\n            adv, grad = masked_step(x, adv, y, target, mask, 0.18)\n\n        for step_size in (0.14, 0.10):\n            adv, _ = masked_step(x, adv, y, target, mask, step_size)\n        return adv.detach()\n\n    with torch.enable_grad():\n        for start in range(0, images.size(0), batch_size):\n            end = start + batch_size\n            x = images[start:end]\n            y = labels[start:end]\n\n            with torch.no_grad():\n                clean_logits = model(x)\n                order = clean_logits.argsort(dim=1, descending=True)\n                alt_targets = order[order.ne(y[:, None])].view(x.size(0), -1)\n                n_targets = min(4 if n_classes <= 10 else 6, alt_targets.size(1))\n                best_margin = decision_margin(clean_logits, y)\n\n            best_adv = x.clone()\n            success = best_margin > 0\n\n            for t in range(n_targets + 1):\n                active = (~success).nonzero(as_tuple=False).squeeze(1)\n                if active.numel() == 0:\n                    break\n                xa, ya = x[active], y[active]\n                target = alt_targets[active, t] if t < n_targets else None\n                cand_adv = attack_once(xa, ya, target)\n\n                with torch.no_grad():\n                    cand_logits = model(cand_adv)\n                    cand_margin = decision_margin(cand_logits, ya)\n\n                improve = cand_margin > best_margin[active]\n                if improve.any():\n                    chosen = active[improve]\n                    best_adv[chosen] = cand_adv[improve]\n                    best_margin[chosen] = cand_margin[improve]\n                success[active] |= cand_margin > 0\n\n            adv_images[start:end] = best_adv\n\n    return adv_images.clamp(0.0, 1.0)\n\n# =====================================================================\n# END EDITABLE REGION\n# =====================================================================",
5  "after_line": 0,
6  "start_line": 3,
7  "end_line": 135
8}

1{
2  "n": 2
3}