attacker_quad.py

import torch
import torch.nn as nn
from torch.optim import Optimizer
from torchvision import transforms

from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Type,
    Union,
)

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal

import dataclasses

import time

import cvxpy as cp
import qpth


def solve_qp(Q, P, G, H):
    B = Q.shape[0]

    if B == 1:
        # Batch size of 1 has weird instabilities
        # I imagine there is a .squeeze() or something inside the QP solver's code
        # that messes up broadcasting dimensions when batch dimension is 1 so let's
        # artificially make 2 solutions when we need 1

        Q = Q.expand(2, -1, -1)
        P = P.expand(2, -1)
        G = G.expand(2, -1, -1)
        H = H.expand(2, -1)

    e = torch.empty(0, device=Q.device)
    z_sol = qpth.qp.QPFunction(verbose=-1, eps=1e-2, check_Q_spd=False)(
        Q, P, G, H, e, e
    )

    if B == 1:
        z_sol = z_sol[:1]

    return z_sol


class QuadAttacK(nn.Module):
    def __init__(
        self,
        pretrained_model: nn.Module,
        optimizer: Optimizer = torch.optim.AdamW,
        optimizer_kwargs: Optional[Dict] = dict(),
        lr: float = 2e-3,
        loss_coefs: Tuple[float, float] = (0.0, 0.0),  # lower, upper bounds
        loss_coefs_decay: float = 0.75,
        binary_search_steps: int = 1,
        max_iterations: int = 30,
        adv_init_coef: float = 0.0,
        margin: float = 0.2,
        warmup_iterations: int = 5,
    ) -> None:
        super().__init__()

        self.model = pretrained_model
        self.optim = optimizer
        self.optim_kwargs = optimizer_kwargs
        # self.optim_kwargs.update(dict(weight_decay=0.0))
        self.lr = lr
        self.loss_coefs = loss_coefs
        self.loss_coefs_decay = loss_coefs_decay
        self.binary_search_steps = binary_search_steps
        self.max_iterations = max_iterations
        self.adv_init_coef = adv_init_coef
        self.margin = margin
        self.warmup_iterations = warmup_iterations

        for p in self.model.parameters():
            p.requires_grad_(False)

        self.model.eval()

    @torch.enable_grad()
    def loss(
        self,
        feats: torch.Tensor,
        attack_targets: torch.Tensor,
        head_W: torch.Tensor,
        head_bias: torch.Tensor,
        device: torch.device,
    ) -> torch.Tensor:
        B, K = attack_targets.shape

        num_feats = head_W.shape[1]
        num_classes = head_W.shape[0]

        # Start with all classes should be less than smallest attack target
        D = -torch.eye(num_classes, device=device)[None].repeat(B, 1, 1)  # [B, C, C]
        attack_targets_write = attack_targets[:, -1][:, None, None].expand(
            -1, D.shape[1], -1
        )
        D.scatter_(
            dim=2,
            index=attack_targets_write,
            src=torch.ones(attack_targets_write.shape, device=device),
        )

        # Clear out the constraint row for each item in the attack targets
        attack_targets_clear = attack_targets[:, :, None].expand(-1, -1, D.shape[-1])
        D.scatter_(
            dim=1,
            index=attack_targets_clear,
            src=torch.zeros(attack_targets_clear.shape, device=device),
        )

        batch_inds = torch.arange(B, device=device)[:, None].expand(-1, K - 1)
        attack_targets_pos = attack_targets[:, :-1]  # [B, K-1]
        attack_targets_neg = attack_targets[:, 1:]  # [B, K-1]

        attack_targets_neg_inds = torch.stack(
            (batch_inds, attack_targets_neg, attack_targets_neg), dim=0
        )  # [3, B, K - 1]
        attack_targets_neg_inds = attack_targets_neg_inds.view(3, -1)

        D[
            attack_targets_neg_inds[0],
            attack_targets_neg_inds[1],
            attack_targets_neg_inds[2],
        ] = -1

        attack_targets_pos_inds = torch.stack(
            (batch_inds, attack_targets_neg, attack_targets_pos), dim=0
        )  # [3, B, K - 1]

        D[
            attack_targets_pos_inds[0],
            attack_targets_pos_inds[1],
            attack_targets_pos_inds[2],
        ] = 1

        A = head_W
        b = head_bias

        Q = 2 * torch.eye(num_feats, device=device)[None].expand(B, -1, -1)

        # We want the solution features to be as close as possible
        # to the current features but also head on the direction of
        # the smallest possible perturbation from the initial predicted
        # features
        anchor_feats = feats

        P = -2 * anchor_feats.expand(B, -1)

        G = -D @ A
        H = -(self.margin - D @ b)

        # Constraints are indexed by smaller logit
        # First attack target isn't smaller than any logit, so its
        # constraint index is redundant, but we keep it for easier parallelization
        # Make this constraint all 0s
        zero_inds = attack_targets[:, 0:1]  # [B, 1]
        H.scatter_(
            dim=1, index=zero_inds, src=torch.zeros(zero_inds.shape, device=device)
        )

        # Clear out unnecessary constraint
        keep_mask = torch.ones_like(H, dtype=torch.bool)
        keep_mask.scatter_(
            dim=1,
            index=zero_inds,
            src=torch.zeros(zero_inds.shape, dtype=torch.bool, device=H.device),
        )

        G = G[keep_mask]
        H = H[keep_mask]

        G = G.view(B, num_classes - 1, num_feats)
        H = H.view(B, num_classes - 1)

        z_sol = solve_qp(Q, P, G, H)

        loss_adv = (feats - z_sol).square().sum(dim=-1)

        return loss_adv

    @torch.enable_grad()
    def attack(
        self,
        x: torch.Tensor,
        gt_labels: torch.Tensor,
        attack_targets: torch.Tensor,
        loss_coefs: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        For a given loss coefficient per image, this function computes
        best-energy attack in the given number of iterations and configuration

        x: [B, C, H, W] [0-1 for colors], not normalized yet
        attack_targets: [B, K] (long)
        gt_labels: [B] (long)
        loss_coefs: [B] (floats)
        """

        loss_coefs = loss_coefs.clone()
        B, K = attack_targets.shape
        device = x.device

        # learnable parameters
        adv_perturbations = nn.Parameter(
            torch.randn(x.shape, device=device, dtype=x.dtype) * self.adv_init_coef,
            requires_grad=True,
        )
        if self.adv_init_coef > 0.0:
            with torch.no_grad():
                x_perturbed = x + adv_perturbations
                x_perturbed = x_perturbed.clamp_(min=0.0, max=1.0)

                adv_perturbations.data = x_perturbed - x

        optim = self.optim([adv_perturbations], lr=self.lr, **self.optim_kwargs)

        best_perturbations = torch.zeros_like(x)
        best_perturbed = torch.zeros_like(x)
        adv_logits = torch.zeros(
            (B, self.model.num_classes()), dtype=x.dtype, device=device
        )
        has_successful_attack = torch.zeros(B, dtype=torch.bool, device=device)  # [B]
        min_energy = torch.full((B,), float("inf"), device=device, dtype=x.dtype)

        head_W, head_bias = self.model.head_matrices()

        for i in range(self.max_iterations):
            if i == self.warmup_iterations:
                # reset optim state
                optim = self.optim([adv_perturbations], lr=self.lr, **self.optim_kwargs)

            x_perturbed = x + adv_perturbations

            cur_logits, cur_feats = self.model(x_perturbed, return_features=True)

            cur_pred_classes = cur_logits.argsort(dim=-1, descending=True)  # [B C]
            attack_successful = (cur_pred_classes[:, :K] == attack_targets).all(
                dim=-1
            )  # [B]
            attack_energy = adv_perturbations.flatten(1).norm(dim=-1)  # [B]

            # print(cur_pred_classes[:, :K])
            # print(attack_targets)

            attack_improved = attack_successful & (attack_energy <= min_energy)

            best_perturbations[attack_improved] = adv_perturbations[
                attack_improved
            ].detach()
            best_perturbed[attack_improved] = x_perturbed[attack_improved].detach()
            if i == 0:
                adv_logits = cur_logits.detach().clone()
            else:
                adv_logits[attack_improved] = cur_logits[attack_improved].detach()
            has_successful_attack[attack_improved] = True
            min_energy[attack_improved] = attack_energy[attack_improved].detach()

            loss_adv = self.loss(
                feats=cur_feats,
                attack_targets=attack_targets,
                head_W=head_W,
                head_bias=head_bias,
                device=device,
            )

            loss_energy = adv_perturbations.flatten(1).square().sum(dim=-1)

            adv_weights = torch.where(
                loss_coefs > 0, loss_coefs, torch.ones_like(loss_coefs)
            )
            energy_weights = torch.where(
                loss_coefs <= 0, -loss_coefs, torch.ones_like(loss_coefs)
            )

            loss = loss_adv * adv_weights + loss_energy * energy_weights
            loss = loss.sum()
            # print(loss)

            optim.zero_grad()
            loss.backward()
            optim.step()

            # If we were successfull let's start taking the norm down
            loss_coefs[attack_improved] *= self.loss_coefs_decay

            # Project perturbation to be within image limits
            with torch.no_grad():
                x_perturbed = x.to(torch.float32) + adv_perturbations
                x_perturbed = x_perturbed.clamp_(min=0.0, max=1.0)

                adv_perturbations.data = x_perturbed - x

                if i == self.max_iterations - 1:
                    x_perturbed = x + adv_perturbations
                    cur_logits = self.model(x_perturbed)[0]

                    cur_pred_classes = cur_logits.argsort(
                        dim=-1, descending=True
                    )  # [B C]
                    attack_successful = (cur_pred_classes[:, :K] == attack_targets).all(
                        dim=-1
                    )  # [B]
                    attack_energy = adv_perturbations.flatten(1).norm(dim=-1)  # [B]

                    attack_improved = attack_successful & (attack_energy <= min_energy)

                    best_perturbations[attack_improved] = adv_perturbations[
                        attack_improved
                    ].detach()
                    best_perturbed[attack_improved] = x_perturbed[
                        attack_improved
                    ].detach()
                    adv_logits[attack_improved] = cur_logits[attack_improved].detach()
                    min_energy[attack_improved] = attack_energy[
                        attack_improved
                    ].detach()

        # adv_logits = self.model(x + best_perturbations)[0]

        return adv_logits, best_perturbed, best_perturbations, min_energy

    def forward(
        self, x: torch.Tensor, gt_labels: torch.Tensor, attack_targets: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        B = x.size(0)
        device = x.device
        loss_coefs_lower = torch.full(
            (B,), fill_value=self.loss_coefs[0], device=device, dtype=x.dtype
        )

        loss_coefs_upper = torch.full(
            (B,), fill_value=self.loss_coefs[1], device=device, dtype=x.dtype
        )

        best_adv_logits = torch.zeros(
            (B, self.model.num_classes()), dtype=x.dtype, device=device
        )
        best_perturbations = torch.zeros_like(x)  # [B, 3, H, W]
        best_perturbed = torch.zeros_like(x)  # [B, 3, H, W]
        best_energy = torch.full((B,), float("inf"), device=device)  # [B]

        for search_step_i in range(self.binary_search_steps):
            # if device.index is None or device.index == 0:
            #     print ("Running binary search step", search_step_i + 1)

            current_loss_coefs = (loss_coefs_lower + loss_coefs_upper) / 2
            current_logits, current_perturbed, current_perturbations, current_energy = (
                self.attack(x, gt_labels, attack_targets, current_loss_coefs)
            )

            current_attack_suceeded = ~torch.isinf(current_energy)

            update_mask = current_energy < best_energy

            if search_step_i == 0:
                best_adv_logits = current_logits.clone()
                best_perturbations = current_perturbations.clone()
                best_perturbed = current_perturbed.clone()
                best_energy = current_energy.clone()
            else:
                best_adv_logits[update_mask] = current_logits[update_mask]
                best_perturbations[update_mask] = current_perturbations[update_mask]
                best_perturbed[update_mask] = current_perturbed[update_mask]
                best_energy[update_mask] = current_energy[update_mask]

            # If we fail to attack, we must increase our loss coef
            loss_coefs_lower[~current_attack_suceeded] = current_loss_coefs[
                ~current_attack_suceeded
            ]

            # If we succeed, we should lower to seek a more frugal attack
            loss_coefs_upper[current_attack_suceeded] = current_loss_coefs[
                current_attack_suceeded
            ]

        return best_adv_logits, best_perturbed, best_perturbations, best_energy