code/models/interpretation.py

import pytorch_lightning as pl
import torch
import torch.nn.functional as F

from .gates import DiffMaskGateInput
from argparse import ArgumentParser
from math import sqrt
from pytorch_lightning.core.optimizer import LightningOptimizer
from torch import Tensor
from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler
from transformers import (
    get_constant_schedule_with_warmup,
    get_constant_schedule,
    ViTForImageClassification,
)
from transformers.models.vit.configuration_vit import ViTConfig
from typing import Optional, Union
from utils.getters_setters import vit_getter, vit_setter
from utils.metrics import accuracy_precision_recall_f1
from utils.optimizer import LookaheadAdam


class ImageInterpretationNet(pl.LightningModule):
    @staticmethod
    def add_model_specific_args(parent_parser: ArgumentParser) -> ArgumentParser:
        parser = parent_parser.add_argument_group("Vision DiffMask")
        parser.add_argument(
            "--alpha",
            type=float,
            default=20.0,
            help="Initial value for the Lagrangian",
        )
        parser.add_argument(
            "--lr",
            type=float,
            default=2e-5,
            help="Learning rate for DiffMask.",
        )
        parser.add_argument(
            "--eps",
            type=float,
            default=0.1,
            help="KL divergence tolerance.",
        )
        parser.add_argument(
            "--no_placeholder",
            action="store_true",
            help="Whether to not use placeholder",
        )
        parser.add_argument(
            "--lr_placeholder",
            type=float,
            default=1e-3,
            help="Learning for mask vectors.",
        )
        parser.add_argument(
            "--lr_alpha",
            type=float,
            default=0.3,
            help="Learning rate for lagrangian optimizer.",
        )
        parser.add_argument(
            "--mul_activation",
            type=float,
            default=15.0,
            help="Value to multiply gate activations.",
        )
        parser.add_argument(
            "--add_activation",
            type=float,
            default=8.0,
            help="Value to add to gate activations.",
        )
        parser.add_argument(
            "--weighted_layer_distribution",
            action="store_true",
            help="Whether to use a weighted distribution when picking a layer in DiffMask forward.",
        )
        return parent_parser

    # Declare variables that will be initialized later
    model: ViTForImageClassification

    def __init__(
        self,
        model_cfg: ViTConfig,
        alpha: float = 1,
        lr: float = 3e-4,
        eps: float = 0.1,
        eps_valid: float = 0.8,
        acc_valid: float = 0.75,
        lr_placeholder: float = 1e-3,
        lr_alpha: float = 0.3,
        mul_activation: float = 10.0,
        add_activation: float = 5.0,
        placeholder: bool = True,
        weighted_layer_pred: bool = False,
    ):
        """A PyTorch Lightning Module for the VisionDiffMask model on the Vision Transformer.

        Args:
            model_cfg (ViTConfig): the configuration of the Vision Transformer model
            alpha (float): the initial value for the Lagrangian
            lr (float): the learning rate for the DiffMask gates
            eps (float): the tolerance for the KL divergence
            eps_valid (float): the tolerance for the KL divergence in the validation step
            acc_valid (float): the accuracy threshold for the validation step
            lr_placeholder (float): the learning rate for the learnable masking embeddings
            lr_alpha (float): the learning rate for the Lagrangian
            mul_activation (float): the value to multiply the gate activations by
            add_activation (float): the value to add to the gate activations
            placeholder (bool): whether to use placeholder embeddings or a zero vector
            weighted_layer_pred (bool): whether to use a weighted distribution when picking a layer
        """
        super().__init__()

        # Save the hyperparameters
        self.save_hyperparameters()

        # Create DiffMask instance
        self.gate = DiffMaskGateInput(
            hidden_size=model_cfg.hidden_size,
            hidden_attention=model_cfg.hidden_size // 4,
            num_hidden_layers=model_cfg.num_hidden_layers + 2,
            max_position_embeddings=1,
            mul_activation=mul_activation,
            add_activation=add_activation,
            placeholder=placeholder,
        )

        # Create the Lagrangian values for the dual optimization
        self.alpha = torch.nn.ParameterList(
            [
                torch.nn.Parameter(torch.ones(()) * alpha)
                for _ in range(model_cfg.num_hidden_layers + 2)
            ]
        )

        # Register buffers for running metrics
        self.register_buffer(
            "running_acc", torch.ones((model_cfg.num_hidden_layers + 2,))
        )
        self.register_buffer(
            "running_l0", torch.ones((model_cfg.num_hidden_layers + 2,))
        )
        self.register_buffer(
            "running_steps", torch.zeros((model_cfg.num_hidden_layers + 2,))
        )

    def set_vision_transformer(self, model: ViTForImageClassification):
        """Set the Vision Transformer model to be used with this module."""
        # Save the model instance as a class attribute
        self.model = model
        # Freeze the model's parameters
        for param in self.model.parameters():
            param.requires_grad = False

    def forward_explainer(
        self, x: Tensor, attribution: bool = False
    ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, int, int]:
        """Performs a forward pass through the explainer (VisionDiffMask) model."""
        # Get the original logits and hidden states from the model
        logits_orig, hidden_states = vit_getter(self.model, x)

        # Add [CLS] token to deal with shape mismatch in self.gate() call
        patch_embeddings = hidden_states[0]
        batch_size = len(patch_embeddings)
        cls_tokens = self.model.vit.embeddings.cls_token.expand(batch_size, -1, -1)
        hidden_states[0] = torch.cat((cls_tokens, patch_embeddings), dim=1)

        # Select the layer to generate the mask from in this pass
        n_hidden = len(hidden_states)
        if self.hparams.weighted_layer_pred:
            # If weighted layer prediction is enabled, use a weighted distribution
            # instead of uniformly picking a layer after a number of steps
            low_weight = (
                lambda i: self.running_acc[i] > 0.75
                and self.running_l0[i] < 0.1
                and self.running_steps[i] > 100
            )
            layers = torch.tensor(list(range(n_hidden)))
            p = torch.tensor([0.1 if low_weight(i) else 1 for i in range(n_hidden)])
            p = p / p.sum()
            idx = p.multinomial(num_samples=1)
            layer_pred = layers[idx].item()
        else:
            layer_pred = torch.randint(n_hidden, ()).item()

        # Set the layer to drop to 0, since we are only interested in masking the input
        layer_drop = 0

        (
            new_hidden_state,
            gates,
            expected_L0,
            gates_full,
            expected_L0_full,
        ) = self.gate(
            hidden_states=hidden_states,
            layer_pred=None
            if attribution
            else layer_pred,  # if attribution, we get all the hidden states
        )

        # Create the list of the new hidden states for the new forward pass
        new_hidden_states = (
            [None] * layer_drop
            + [new_hidden_state]
            + [None] * (n_hidden - layer_drop - 1)
        )

        # Get the new logits from the masked input
        logits, _ = vit_setter(self.model, x, new_hidden_states)

        return (
            logits,
            logits_orig,
            gates,
            expected_L0,
            gates_full,
            expected_L0_full,
            layer_drop,
            layer_pred,
        )

    def get_mask(self, x: Tensor,
                 idx: int = -1,
                 aggregated_mask: bool = True,
                 ) -> dict[str, Tensor]:
        """
        Generates a mask for the given input.
        Args:
            x: the input to generate the mask for
            idx: the index of the layer to generate the mask from
            aggregated_mask: whether to use an aggregative mask from each layer
        Returns:
            a dictionary containing the mask, kl divergence and the predicted class
        """

        # Pass from forward explainer with attribution=True
        (
            logits,
            logits_orig,
            gates,
            expected_L0,
            gates_full,
            expected_L0_full,
            layer_drop,
            layer_pred,
        ) = self.forward_explainer(x, attribution=True)

        # Calculate KL-divergence
        kl_div = torch.distributions.kl_divergence(
            torch.distributions.Categorical(logits=logits_orig),
            torch.distributions.Categorical(logits=logits),
        )

        # Get predicted class
        pred_class = logits.argmax(-1)

        # Calculate mask
        if aggregated_mask:
            mask = expected_L0_full[:, :, idx].exp()
        else:
            mask = gates_full[:, :, idx]

        mask = mask[:, 1:]

        C, H, W = x.shape[1:]  # channels, height, width
        B, P = mask.shape  # batch, patches
        N = int(sqrt(P))  # patches per side
        S = int(H / N)  # patch size

        # Reshape mask to match input shape
        mask = mask.reshape(B, 1, N, N)
        mask = F.interpolate(mask, scale_factor=S)
        mask = mask.reshape(B, H, W)

        return {"mask": mask, "kl_div": kl_div, "pred_class": pred_class,
                "logits": logits, "logits_orig": logits_orig}

    def forward(self, x: Tensor) -> Tensor:
        return self.model(x).logits

    def training_step(self, batch: tuple[Tensor, Tensor], *args, **kwargs) -> dict:
        # Unpack the batch
        x, y = batch

        # Pass the batch through the explainer (VisionDiffMask) model
        (
            logits,
            logits_orig,
            gates,
            expected_L0,
            gates_full,
            expected_L0_full,
            layer_drop,
            layer_pred,
        ) = self.forward_explainer(x)

        # Calculate the KL-divergence loss term
        loss_c = (
            torch.distributions.kl_divergence(
                torch.distributions.Categorical(logits=logits_orig),
                torch.distributions.Categorical(logits=logits),
            )
            - self.hparams.eps
        )

        # Calculate the L0 loss term
        loss_g = expected_L0.mean(-1)

        # Calculate the full loss term
        loss = self.alpha[layer_pred] * loss_c + loss_g

        # Calculate the accuracy
        acc, _, _, _ = accuracy_precision_recall_f1(
            logits.argmax(-1), logits_orig.argmax(-1), average=True
        )

        # Calculate the average L0 loss
        l0 = expected_L0.exp().mean(-1)

        outputs_dict = {
            "loss_c": loss_c.mean(-1),
            "loss_g": loss_g.mean(-1),
            "alpha": self.alpha[layer_pred].mean(-1),
            "acc": acc,
            "l0": l0.mean(-1),
            "layer_pred": layer_pred,
            "r_acc": self.running_acc[layer_pred],
            "r_l0": self.running_l0[layer_pred],
            "r_steps": self.running_steps[layer_pred],
            "debug_loss": loss.mean(-1),
        }

        outputs_dict = {
            "loss": loss.mean(-1),
            **outputs_dict,
            "log": outputs_dict,
            "progress_bar": outputs_dict,
        }

        self.log(
            "loss", outputs_dict["loss"], on_step=True, on_epoch=True, prog_bar=True
        )
        self.log(
            "loss_c", outputs_dict["loss_c"], on_step=True, on_epoch=True, prog_bar=True
        )
        self.log(
            "loss_g", outputs_dict["loss_g"], on_step=True, on_epoch=True, prog_bar=True
        )
        self.log("acc", outputs_dict["acc"], on_step=True, on_epoch=True, prog_bar=True)
        self.log("l0", outputs_dict["l0"], on_step=True, on_epoch=True, prog_bar=True)
        self.log(
            "alpha", outputs_dict["alpha"], on_step=True, on_epoch=True, prog_bar=True
        )

        outputs_dict = {
            "{}{}".format("" if self.training else "val_", k): v
            for k, v in outputs_dict.items()
        }

        if self.training:
            self.running_acc[layer_pred] = (
                self.running_acc[layer_pred] * 0.9 + acc * 0.1
            )
            self.running_l0[layer_pred] = (
                self.running_l0[layer_pred] * 0.9 + l0.mean(-1) * 0.1
            )
            self.running_steps[layer_pred] += 1

        return outputs_dict

    def validation_epoch_end(self, outputs: list[dict]):
        outputs_dict = {
            k: [e[k] for e in outputs if k in e]
            for k in ("val_loss_c", "val_loss_g", "val_acc", "val_l0")
        }

        outputs_dict = {k: sum(v) / len(v) for k, v in outputs_dict.items()}

        outputs_dict["val_loss_c"] += self.hparams.eps

        outputs_dict = {
            "val_loss": outputs_dict["val_l0"]
            if outputs_dict["val_loss_c"] <= self.hparams.eps_valid
            and outputs_dict["val_acc"] >= self.hparams.acc_valid
            else torch.full_like(outputs_dict["val_l0"], float("inf")),
            **outputs_dict,
            "log": outputs_dict,
        }

        return outputs_dict

    def configure_optimizers(self) -> tuple[list[Optimizer], list[_LRScheduler]]:
        optimizers = [
            LookaheadAdam(
                params=[
                    {
                        "params": self.gate.g_hat.parameters(),
                        "lr": self.hparams.lr,
                    },
                    {
                        "params": self.gate.placeholder.parameters()
                        if isinstance(self.gate.placeholder, torch.nn.ParameterList)
                        else [self.gate.placeholder],
                        "lr": self.hparams.lr_placeholder,
                    },
                ],
                # centered=True, # this is for LookaheadRMSprop
            ),
            LookaheadAdam(
                params=[self.alpha]
                if isinstance(self.alpha, torch.Tensor)
                else self.alpha.parameters(),
                lr=self.hparams.lr_alpha,
            ),
        ]

        schedulers = [
            {
                "scheduler": get_constant_schedule_with_warmup(optimizers[0], 12 * 100),
                "interval": "step",
            },
            get_constant_schedule(optimizers[1]),
        ]
        return optimizers, schedulers

    def optimizer_step(
        self,
        epoch: int,
        batch_idx: int,
        optimizer: Union[Optimizer, LightningOptimizer],
        optimizer_idx: int = 0,
        optimizer_closure: Optional[callable] = None,
        on_tpu: bool = False,
        using_native_amp: bool = False,
        using_lbfgs: bool = False,
    ):
        # Optimizer 0: Minimize loss w.r.t. DiffMask's parameters
        if optimizer_idx == 0:
            # Gradient ascent on the model's parameters
            optimizer.step(closure=optimizer_closure)
            optimizer.zero_grad()
            for g in optimizer.param_groups:
                for p in g["params"]:
                    p.grad = None

        # Optimizer 1: Maximize loss w.r.t. the Langrangian
        elif optimizer_idx == 1:
            # Reverse the sign of the Langrangian's gradients
            for i in range(len(self.alpha)):
                if self.alpha[i].grad:
                    self.alpha[i].grad *= -1

            # Gradient ascent on the Langrangian
            optimizer.step(closure=optimizer_closure)
            optimizer.zero_grad()
            for g in optimizer.param_groups:
                for p in g["params"]:
                    p.grad = None

            # Clip the Lagrangian's values
            for i in range(len(self.alpha)):
                self.alpha[i].data = torch.where(
                    self.alpha[i].data < 0,
                    torch.full_like(self.alpha[i].data, 0),
                    self.alpha[i].data,
                )
                self.alpha[i].data = torch.where(
                    self.alpha[i].data > 200,
                    torch.full_like(self.alpha[i].data, 200),
                    self.alpha[i].data,
                )

    def on_save_checkpoint(self, ckpt: dict):
        # Remove VIT from checkpoint as we can load it dynamically
        keys = list(ckpt["state_dict"].keys())
        for key in keys:
            if key.startswith("model."):
                del ckpt["state_dict"][key]