model_under_attack.py

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 torch
from torch import nn
from torch.func import functional_call, jacrev, vmap
from torch.cuda.amp import autocast

import timm
from timm.models import create_model

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)


def filter_logits(logits, remove_logits):
    keep_mask = torch.ones(logits.shape, device=logits.device, dtype=torch.bool)
    false_fill = torch.zeros(
        remove_logits.shape, device=logits.device, dtype=torch.bool
    )
    keep_mask = torch.scatter(keep_mask, dim=-1, index=remove_logits, src=false_fill)

    num_keep_logits = logits.shape[1] - remove_logits.shape[1]

    logits = logits[keep_mask]
    logits = logits.view(remove_logits.shape[0], num_keep_logits)

    return logits


def scaled_softmax(x, p):
    x_scaled = x[..., None] * p
    x_scaled = x_scaled.softmax(dim=-2)

    x_scaled = x[..., None] * x_scaled
    x_scaled = x_scaled.sum(dim=-2)

    return x_scaled


class BaseModel(nn.Module):
    def __init__(
        self,
        model_name: str,
        data_mean: Sequence[Union[float, int]] = IMAGENET_DEFAULT_MEAN,
        data_std: Sequence[Union[float, int]] = IMAGENET_DEFAULT_STD,
        logits_buffer_size_multiplier: int = 5,
    ) -> None:
        super().__init__()
        self.model_name = model_name
        self.logits_buffer_size_multiplier = logits_buffer_size_multiplier

        assert len(data_mean) == 3, (
            "`mean` should have 3 values, to be compatible with "
            f"RGB image, but got {len(data_mean)} values"
        )
        assert len(data_std) == 3, (
            "`std` should have 3 values, to be compatible with RGB "
            f"image, but got {len(data_std)} values"
        )
        self.register_buffer("mean", torch.tensor(data_mean).view(-1, 1, 1), False)
        self.register_buffer("std", torch.tensor(data_std).view(-1, 1, 1), False)

    @property
    def head(self):
        raise NotImplementedError("Not implemented for base class")

    def head_features(self):
        return self.head.in_features

    def num_classes(self):
        return self.head.out_features

    def head_matrices(self):
        return self.head.weight, self.head.bias

    def reorder_and_reduce_logits(
        self, logits: torch.Tensor, attack_targets: torch.Tensor
    ):
        """
        Project logits down to a subspace from the logit covariance
        logits: [..., C]
        attack_targets: [..., topK] logits to maintain exactly in subspace projection
        """
        C = logits.shape[-1]
        topK = attack_targets.shape[-1]

        keep_mask = torch.ones(logits.shape, device=logits.device, dtype=torch.bool)
        false_fill = torch.zeros(
            attack_targets.shape, device=logits.device, dtype=torch.bool
        )
        keep_mask = torch.scatter(
            keep_mask, dim=-1, index=attack_targets, src=false_fill
        )

        # Doing this mess because we are not allowed to use dynamic shape mask ops in VMAP
        num_non_attack_targets = C - topK
        non_exact_scatter_inds = torch.full_like(
            logits, dtype=torch.long, fill_value=num_non_attack_targets
        )

        num_dummy_dims = logits.dim() - 1
        scatter_fill_vals = torch.arange(num_non_attack_targets, device=logits.device)[
            [None] * num_dummy_dims
        ]
        scatter_fill_vals = scatter_fill_vals.expand(list(logits.shape[:-1]) + [-1])

        try:
            non_exact_scatter_inds.masked_scatter_(keep_mask, scatter_fill_vals)
        except:
            print(
                attack_targets,
                keep_mask.shape,
                scatter_fill_vals.shape,
                non_exact_scatter_inds.shape,
            )

        # ## Index the rows and column positions to assign into
        # row_ids, col_ids = torch.where(keep_mask)  # [N], [N]

        # # Flatten and assign directly (this is vmap-safe)
        # non_exact_scatter_inds[row_ids, col_ids] = scatter_fill_vals.view(-1)

        non_exact_logit_vals = torch.empty(
            list(logits.shape[:-1]) + [num_non_attack_targets + 1], device=logits.device
        )
        non_exact_logit_vals = non_exact_logit_vals.scatter(
            dim=-1, index=non_exact_scatter_inds, src=logits
        )

        # Dump the dummy vals that contain some of the exact logits
        non_exact_logit_vals = non_exact_logit_vals[..., :-1]

        exact_logit_vals = torch.gather(
            logits, dim=1, index=attack_targets
        )  # [..., topK]

        non_exact_logit_vals = non_exact_logit_vals.sort(dim=-1, descending=True).values

        keep_non_exact = topK * self.logits_buffer_size_multiplier
        smooth_max_logits = scaled_softmax(
            non_exact_logit_vals[..., keep_non_exact:],
            p=torch.tensor([1 / 2, 1.0, 999.0], device=logits.device),
        )

        return torch.cat(
            [
                exact_logit_vals,
                non_exact_logit_vals[..., :keep_non_exact],
                smooth_max_logits,
            ],
            dim=-1,
        )

    @torch.no_grad
    def compute_logits_to_data_jacobian(
        self, x: torch.Tensor, attack_targets: torch.Tensor, chunk_size: int = 100
    ):

        def _single(x_single: torch.Tensor, attack_targets_single: torch.Tensor):
            x_single = x_single.unsqueeze(0)
            attack_targets_single = attack_targets_single.unsqueeze(0)

            compressed_logits, logits = self(
                x_single, attack_targets=attack_targets_single
            )

            return compressed_logits[0], (compressed_logits[0], logits[0])

        single_jac = jacrev(_single, has_aux=True, chunk_size=chunk_size)
        batch_jac = vmap(single_jac)

        # with autocast(enabled=False):
        jacobian, (compressed_logits, logits) = batch_jac(x, attack_targets)

        return compressed_logits, logits, jacobian

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError("Not implemented for base class")

    def forward(
        self,
        x: torch.Tensor,
        return_features: bool = False,
        attack_targets: Optional[torch.Tensor] = None,
        return_logits: bool = True,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        """
        x: [B, 3 (RGB), H, W] image (float) [0,1] without normalization

        returns: [B, C] class logits + ([B, D] features)
        """
        x = (x - self.mean.to(x.device)) / self.std.to(x.device)
        feats = self.forward_features(x)
        if return_logits or attack_targets is not None:
            logits = self.head(feats)

        output = []
        if attack_targets is not None:
            compressed_logits = self.reorder_and_reduce_logits(logits, attack_targets)
            output.append(compressed_logits)

        if return_logits:
            output.append(logits)

        if return_features:
            output.append(feats)

        return tuple(output)


class TimmModel(BaseModel):
    """
    Wrap models from the timm package
    """

    def __init__(
        self,
        model_name: Optional[str] = None,
        model: Optional[nn.Module] = None,
        data_mean: Sequence[Union[float, int]] = IMAGENET_DEFAULT_MEAN,
        data_std: Sequence[Union[float, int]] = IMAGENET_DEFAULT_STD,
        logits_buffer_size_multiplier: int = 5,
    ) -> None:
        super().__init__(model_name, data_mean, data_std)

        self.logits_buffer_size_multiplier = logits_buffer_size_multiplier

        if model is None:
            assert model_name is not None
            self.model = create_model(self.model_name, pretrained=True)
        else:
            self.model = model

        self.model.eval()

        # if hasattr(self.model, 'pretrained_cfg'):
        #     mean = self.model.pretrained_cfg.get('mean', None)
        #     std = self.model.pretrained_cfg.get('std', None)
        #     if mean is not None and std is not None:
        #         self.mean = torch.tensor(mean).view(-1, 1, 1)
        #         self.std = torch.tensor(std).view(-1, 1, 1)

    @property
    def head(self):
        if hasattr(self.model, "head"):
            if isinstance(self.model.head, torch.nn.Linear):
                return self.model.head
            else:
                return self.model.head.fc
        elif hasattr(self.model, "classifier"):  # e.g. densenet
            return self.model.classifier
        elif hasattr(self.model, "fc"):  # e.g. resnet
            return self.model.fc
        else:
            raise ValueError(f"not found head classifier")

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        feats = self.model.forward_features(x)
        if hasattr(self.model, "global_pool") and callable(
            self.model.global_pool
        ):  # e.g., resnet, densenet
            feats = self.model.global_pool(feats)
        elif hasattr(self.model, "pool") and callable(self.model.pool):  # e.g., vit
            feats = self.model.pool(feats)
            feats = self.model.fc_norm(feats)
            feats = self.model.head_drop(feats)
        else:
            raise ValueError(f"not found pooling layer")  # add more if needed
        return feats