aiops/ContraLSP/abstudy/gatemasknn_no_smooth.py [173:316]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        return mask

    def representation(self, x):
        mask = self.refactor_mask(self.mask, x)
        mask = self.hard_sigmoid(mask)
        return mask.detach().cpu()


class GateMaskNet(Net):
    """
    Gate mask model as a Pytorch Lightning model.

    Args:
        forward_func (callable): The forward function of the model or any
            modification of it.
        preservation_mode (bool): If ``True``, uses the method in
            preservation mode. Otherwise, uses the deletion mode.
            Default to ``True``
        model (nnn.Module): A model used to recreate the original
            predictions, in addition to the mask. Default to ``None``
        batch_size (int): Batch size of the model. Default to 32
        lambda_1 (float): Weighting for the mask loss. Default to 1.
        lambda_2 (float): Weighting for the model output loss. Default to 1.
        loss (str, callable): Which loss to use. Default to ``'mse'``
        optim (str): Which optimizer to use. Default to ``'adam'``
        lr (float): Learning rate. Default to 1e-3
        lr_scheduler (dict, str): Learning rate scheduler. Either a dict
            (custom scheduler) or a string. Default to ``None``
        lr_scheduler_args (dict): Additional args for the scheduler.
            Default to ``None``
        l2 (float): L2 regularisation. Default to 0.0
    """

    def __init__(
        self,
        forward_func: Callable,
        preservation_mode: bool = True,
        model: nn.Module = None,
        batch_size: int = 32,
        lambda_1: float = 1.0,
        lambda_2: float = 1.0,
        factor_dilation=10.0,
        based=0.5,
        loss: Union[str, Callable] = "mse",
        optim: str = "adam",
        lr: float = 0.001,
        lr_scheduler: Union[dict, str] = None,
        lr_scheduler_args: dict = None,
        l2: float = 0.0,
    ):
        mask = GateMaskNN(
            forward_func=forward_func,
            model=model,
            batch_size=batch_size,
            factor_dilation=factor_dilation,
            based=based
        )

        super().__init__(
            layers=mask,
            loss=loss,
            optim=optim,
            lr=lr,
            lr_scheduler=lr_scheduler,
            lr_scheduler_args=lr_scheduler_args,
            l2=l2,
        )

        self.preservation_mode = preservation_mode
        self.lambda_1 = lambda_1
        self.lambda_2 = lambda_2
        self.based = based

    def forward(self, *args, **kwargs) -> th.Tensor:
        return self.net(*args, **kwargs)

    def step(self, batch, batch_idx, stage):
        # x is the data to be perturbed
        # y is the same data without perturbation
        x, y, baselines, target, *additional_forward_args = batch

        # If additional_forward_args is only one None,
        # set it to None
        if additional_forward_args == [None]:
            additional_forward_args = None

        # Get perturbed output
        # y_hat1 is computed by masking important features
        # y_hat2 is computed by masking unimportant features
        if additional_forward_args is None:
            y_hat1, y_hat2 = self(x.float(), batch_idx, baselines, target)
        else:
            y_hat1, y_hat2 = self(
                x.float(),
                batch_idx,
                baselines,
                target,
                *additional_forward_args,
            )

        # Get unperturbed output for inputs and baselines
        y_target1 = _run_forward(
            forward_func=self.net.forward_func,
            inputs=y,
            target=target,
            additional_forward_args=tuple(additional_forward_args)
            if additional_forward_args is not None
            else None,
        )

        # Add L1 loss
        mask_ = self.net.mask[
                self.net.batch_size
                * batch_idx: self.net.batch_size * (batch_idx + 1)
                ]
        reg = 0.5 + 0.5 * th.erf(self.net.refactor_mask(mask_, x) / (self.net.sigma * np.sqrt(2)))
        # reg = self.net.refactor_mask(mask_, x).abs()

        # trend_reg = self.net.trend_info(x).abs().mean()
        mask_loss = self.lambda_1 * reg.mean()
        # mask_loss = self.lambda_1 * th.sum(reg, dim=[1,2]).mean()

        triplet_loss = 0
        if self.net.model is not None:
            condition = self.net.model(x - baselines)
            if self.net.batch_size > 1:
                triplet_loss = self.lambda_2 * self._triplet_loss(condition)
            else:
                triplet_loss = self.lambda_2 * condition.abs().mean()

        # Add preservation and deletion losses if required
        if self.preservation_mode:
            main_loss = self.loss(y_hat1, y_target1)
        else:
            main_loss = -1. * self.loss(y_hat2, y_target1)

        loss = main_loss + mask_loss + triplet_loss

        # test log
        _test_mask = self.net.representation(x)
        test = (_test_mask[_test_mask > 0]).sum()
        lambda_1_t = self.lambda_1
        lambda_2_t = self.lambda_2
        reg_t = reg.mean()
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



aiops/ContraLSP/attribution/gatemasknn.py [191:334]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        return mask

    def representation(self, x):
        mask = self.refactor_mask(self.mask, x)
        mask = self.hard_sigmoid(mask)
        return mask.detach().cpu()


class GateMaskNet(Net):
    """
    Gate mask model as a Pytorch Lightning model.

    Args:
        forward_func (callable): The forward function of the model or any
            modification of it.
        preservation_mode (bool): If ``True``, uses the method in
            preservation mode. Otherwise, uses the deletion mode.
            Default to ``True``
        model (nnn.Module): A model used to recreate the original
            predictions, in addition to the mask. Default to ``None``
        batch_size (int): Batch size of the model. Default to 32
        lambda_1 (float): Weighting for the mask loss. Default to 1.
        lambda_2 (float): Weighting for the model output loss. Default to 1.
        loss (str, callable): Which loss to use. Default to ``'mse'``
        optim (str): Which optimizer to use. Default to ``'adam'``
        lr (float): Learning rate. Default to 1e-3
        lr_scheduler (dict, str): Learning rate scheduler. Either a dict
            (custom scheduler) or a string. Default to ``None``
        lr_scheduler_args (dict): Additional args for the scheduler.
            Default to ``None``
        l2 (float): L2 regularisation. Default to 0.0
    """

    def __init__(
        self,
        forward_func: Callable,
        preservation_mode: bool = True,
        model: nn.Module = None,
        batch_size: int = 32,
        lambda_1: float = 1.0,
        lambda_2: float = 1.0,
        factor_dilation=10.0,
        based=0.5,
        loss: Union[str, Callable] = "mse",
        optim: str = "adam",
        lr: float = 0.001,
        lr_scheduler: Union[dict, str] = None,
        lr_scheduler_args: dict = None,
        l2: float = 0.0,
    ):
        mask = GateMaskNN(
            forward_func=forward_func,
            model=model,
            batch_size=batch_size,
            factor_dilation=factor_dilation,
            based=based
        )

        super().__init__(
            layers=mask,
            loss=loss,
            optim=optim,
            lr=lr,
            lr_scheduler=lr_scheduler,
            lr_scheduler_args=lr_scheduler_args,
            l2=l2,
        )

        self.preservation_mode = preservation_mode
        self.lambda_1 = lambda_1
        self.lambda_2 = lambda_2
        self.based = based

    def forward(self, *args, **kwargs) -> th.Tensor:
        return self.net(*args, **kwargs)

    def step(self, batch, batch_idx, stage):
        # x is the data to be perturbed
        # y is the same data without perturbation
        x, y, baselines, target, *additional_forward_args = batch

        # If additional_forward_args is only one None,
        # set it to None
        if additional_forward_args == [None]:
            additional_forward_args = None

        # Get perturbed output
        # y_hat1 is computed by masking important features
        # y_hat2 is computed by masking unimportant features
        if additional_forward_args is None:
            y_hat1, y_hat2 = self(x.float(), batch_idx, baselines, target)
        else:
            y_hat1, y_hat2 = self(
                x.float(),
                batch_idx,
                baselines,
                target,
                *additional_forward_args,
            )

        # Get unperturbed output for inputs and baselines
        y_target1 = _run_forward(
            forward_func=self.net.forward_func,
            inputs=y,
            target=target,
            additional_forward_args=tuple(additional_forward_args)
            if additional_forward_args is not None
            else None,
        )

        # Add L1 loss
        mask_ = self.net.mask[
                self.net.batch_size
                * batch_idx: self.net.batch_size * (batch_idx + 1)
                ]
        reg = 0.5 + 0.5 * th.erf(self.net.refactor_mask(mask_, x) / (self.net.sigma * np.sqrt(2)))
        # reg = self.net.refactor_mask(mask_, x).abs()

        # trend_reg = self.net.trend_info(x).abs().mean()
        mask_loss = self.lambda_1 * reg.mean()
        # mask_loss = self.lambda_1 * th.sum(reg, dim=[1,2]).mean()

        triplet_loss = 0
        if self.net.model is not None:
            condition = self.net.model(x - baselines)
            if self.net.batch_size > 1:
                triplet_loss = self.lambda_2 * self._triplet_loss(condition)
            else:
                triplet_loss = self.lambda_2 * condition.abs().mean()

        # Add preservation and deletion losses if required
        if self.preservation_mode:
            main_loss = self.loss(y_hat1, y_target1)
        else:
            main_loss = -1. * self.loss(y_hat2, y_target1)

        loss = main_loss + mask_loss + triplet_loss

        # test log
        _test_mask = self.net.representation(x)
        test = (_test_mask[_test_mask > 0]).sum()
        lambda_1_t = self.lambda_1
        lambda_2_t = self.lambda_2
        reg_t = reg.mean()
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -