taps

GradExpander

Bases: Function

A custom autograd function that scales the gradient during the backward pass. This function allows you to define a custom forward and backward pass for a PyTorch operation. The forward pass simply returns the input tensor, while the backward pass scales the gradient by a specified factor alpha. Methods: forward(ctx, x, alpha: float = 1): backward(ctx, grad_x):

Source code in CTRAIN/bound/taps.py

class GradExpander(torch.autograd.Function):
    """
    A custom autograd function that scales the gradient during the backward pass.
    This function allows you to define a custom forward and backward pass for a 
    PyTorch operation. The forward pass simply returns the input tensor, while 
    the backward pass scales the gradient by a specified factor `alpha`.
    Methods:
        forward(ctx, x, alpha: float = 1):
        backward(ctx, grad_x):

    """

    @staticmethod
    def forward(ctx, x, alpha:float=1):
        """
        Forward pass for the custom operation.

        Args:
            ctx: The context object that can be used to stash information
                for backward computation.
            x: The input tensor.
            alpha (float, optional): A scaling factor. Defaults to 1.

        Returns:
            (torch.Tensor): The input tensor `x`.
        """
        ctx.alpha = alpha
        return x

    @staticmethod
    def backward(ctx, grad_x):
        """
        Performs the backward pass for the custom autograd function.

        Args:
            ctx: The context object that can be used to stash information for backward computation.
            grad_x: The gradient of the loss with respect to the output of the forward pass.

        Returns:
            (Tuple[Tensor, None]): A tuple containing the gradient of the loss with respect to the input of the forward pass and None (as there is no gradient with respect to the second input).
        """
        return ctx.alpha * grad_x, None

backward(ctx, grad_x) staticmethod

Performs the backward pass for the custom autograd function.

Parameters:

Name	Type	Description	Default
ctx		The context object that can be used to stash information for backward computation.	required
grad_x		The gradient of the loss with respect to the output of the forward pass.	required

Returns:

Type	Description
Tuple[Tensor, None]	A tuple containing the gradient of the loss with respect to the input of the forward pass and None (as there is no gradient with respect to the second input).

Source code in CTRAIN/bound/taps.py

@staticmethod
def backward(ctx, grad_x):
    """
    Performs the backward pass for the custom autograd function.

    Args:
        ctx: The context object that can be used to stash information for backward computation.
        grad_x: The gradient of the loss with respect to the output of the forward pass.

    Returns:
        (Tuple[Tensor, None]): A tuple containing the gradient of the loss with respect to the input of the forward pass and None (as there is no gradient with respect to the second input).
    """
    return ctx.alpha * grad_x, None

forward(ctx, x, alpha=1) staticmethod

Forward pass for the custom operation.

Parameters:

Name	Type	Description	Default
ctx		The context object that can be used to stash information for backward computation.	required
x		The input tensor.	required
alpha	float	A scaling factor. Defaults to 1.	1

Returns:

Type	Description
Tensor	The input tensor x.

Source code in CTRAIN/bound/taps.py

@staticmethod
def forward(ctx, x, alpha:float=1):
    """
    Forward pass for the custom operation.

    Args:
        ctx: The context object that can be used to stash information
            for backward computation.
        x: The input tensor.
        alpha (float, optional): A scaling factor. Defaults to 1.

    Returns:
        (torch.Tensor): The input tensor `x`.
    """
    ctx.alpha = alpha
    return x

RectifiedLinearGradientLink

Bases: Function

RectifiedLinearGradientLink is a custom autograd function that establishes a rectified linear gradient link between the IBP bounds of the feature extractor (lb, ub) and the PGD bounds (x_adv) of the classifier. This function is not a valid gradient with respect to the forward function.

Attributes:

Name	Type	Description
c	float	A constant used to determine the slope.
tol	float	A tolerance value to avoid division by zero.

Methods:

Name	Description
forward	float, tol: float)
backward

Source code in CTRAIN/bound/taps.py

class RectifiedLinearGradientLink(torch.autograd.Function):
    """
    RectifiedLinearGradientLink is a custom autograd function that establishes a rectified linear gradient link 
    between the IBP bounds of the feature extractor (lb, ub) and the 
    PGD bounds (x_adv) of the classifier. This function is not a valid gradient with respect 
    to the forward function.

    Attributes:
        c (float): A constant used to determine the slope.
        tol (float): A tolerance value to avoid division by zero.

    Methods:
        forward(ctx, lb, ub, x, c: float, tol: float)
        backward(ctx, grad_x):
    """
    @staticmethod
    def forward(ctx, lb, ub, x, c:float, tol:float):
        """
        Saves the input tensors and constants for backward computation.

        Args:
            ctx: Context object to save information for backward computation.
            lb: Lower bound tensor.
            ub: Upper bound tensor.
            x: Input tensor.
            c (float): A constant used to determine the slope.
            tol (float): A tolerance value to avoid division by zero.

        Returns:
            (Tensor): The input tensor x.
        """
        ctx.save_for_backward(lb, ub, x)
        ctx.c = c
        ctx.tol = tol
        return x
    @staticmethod
    def backward(ctx, grad_x):
        """
        Computes the gradient of the loss with respect to the input bounds (lb, ub).

        Args:
            ctx: Context object containing saved tensors and constants.
            grad_x: Gradient of the loss with respect to the output of the forward function.

        Returns:
            (Tuple[Tensor, Tensor, None, None, None]): Gradients with respect to lb, ub, and None for other inputs.
        """
        lb, ub, x = ctx.saved_tensors
        c, tol = ctx.c, ctx.tol
        slackness = c * (ub - lb)
        # handle grad w.r.t. ub
        thre = (ub - slackness)
        Rectifiedgrad_mask = (x >= thre)
        grad_ub = (Rectifiedgrad_mask * grad_x * (x - thre).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
        # handle grad w.r.t. lb
        thre = (lb + slackness)
        Rectifiedgrad_mask = (x <= thre)
        grad_lb = (Rectifiedgrad_mask * grad_x * (thre - x).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
        # we don't need grad w.r.t. x and param
        return grad_lb, grad_ub, None, None, None

backward(ctx, grad_x) staticmethod

Computes the gradient of the loss with respect to the input bounds (lb, ub).

Parameters:

Name	Type	Description	Default
ctx		Context object containing saved tensors and constants.	required
grad_x		Gradient of the loss with respect to the output of the forward function.	required

Returns:

Type	Description
Tuple[Tensor, Tensor, None, None, None]	Gradients with respect to lb, ub, and None for other inputs.

Source code in CTRAIN/bound/taps.py

@staticmethod
def backward(ctx, grad_x):
    """
    Computes the gradient of the loss with respect to the input bounds (lb, ub).

    Args:
        ctx: Context object containing saved tensors and constants.
        grad_x: Gradient of the loss with respect to the output of the forward function.

    Returns:
        (Tuple[Tensor, Tensor, None, None, None]): Gradients with respect to lb, ub, and None for other inputs.
    """
    lb, ub, x = ctx.saved_tensors
    c, tol = ctx.c, ctx.tol
    slackness = c * (ub - lb)
    # handle grad w.r.t. ub
    thre = (ub - slackness)
    Rectifiedgrad_mask = (x >= thre)
    grad_ub = (Rectifiedgrad_mask * grad_x * (x - thre).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
    # handle grad w.r.t. lb
    thre = (lb + slackness)
    Rectifiedgrad_mask = (x <= thre)
    grad_lb = (Rectifiedgrad_mask * grad_x * (thre - x).clamp(min=0.5*tol) / slackness.clamp(min=tol)).sum(dim=0, keepdim=True)
    # we don't need grad w.r.t. x and param
    return grad_lb, grad_ub, None, None, None

forward(ctx, lb, ub, x, c, tol) staticmethod

Saves the input tensors and constants for backward computation.

Parameters:

Name	Type	Description	Default
ctx		Context object to save information for backward computation.	required
lb		Lower bound tensor.	required
ub		Upper bound tensor.	required
x		Input tensor.	required
c	float	A constant used to determine the slope.	required
tol	float	A tolerance value to avoid division by zero.	required

Returns:

Type	Description
Tensor	The input tensor x.

Source code in CTRAIN/bound/taps.py

@staticmethod
def forward(ctx, lb, ub, x, c:float, tol:float):
    """
    Saves the input tensors and constants for backward computation.

    Args:
        ctx: Context object to save information for backward computation.
        lb: Lower bound tensor.
        ub: Upper bound tensor.
        x: Input tensor.
        c (float): A constant used to determine the slope.
        tol (float): A tolerance value to avoid division by zero.

    Returns:
        (Tensor): The input tensor x.
    """
    ctx.save_for_backward(lb, ub, x)
    ctx.c = c
    ctx.tol = tol
    return x

bound_taps(original_model, hardened_model, bounded_blocks, data, target, n_classes, ptb, device='cuda', pgd_steps=20, pgd_restarts=1, pgd_step_size=0.2, pgd_decay_factor=0.2, pgd_decay_checkpoints=(5, 7), gradient_link_thresh=0.5, gradient_link_tolerance=1e-05, propagation='IBP', sabr_args=None)

Compute the bounds of the model's output using the TAPS method.

Parameters:

Name	Type	Description	Default
original_model	Module	The original neural network model.	required
hardened_model	BoundedModule	The auto_LiRPA model.	required
bounded_blocks	list	List of bounded blocks of the model.	required
data	Tensor	The input data tensor.	required
target	Tensor	The target labels tensor.	required
n_classes	int	The number of classes for classification.	required
ptb	PerturbationLpNorm	The perturbation object defining the perturbation set.	required
device	str	The device to run the computation on. Default is 'cuda'.	'cuda'
pgd_steps	int	The number of steps for the PGD attack. Default is 20.	20
pgd_restarts	int	The number of restarts for the PGD attack. Default is 1.	1
pgd_step_size	float	The step size for the PGD attack. Default is 0.2.	0.2
pgd_decay_factor	float	The decay factor for the PGD attack. Default is 0.2.	0.2
pgd_decay_checkpoints	tuple	The decay checkpoints for the PGD attack. Default is (5, 7).	(5, 7)
gradient_link_thresh	float	The threshold for gradient linking. Default is 0.5.	0.5
gradient_link_tolerance	float	The tolerance for gradient linking. Default is 1e-05.	1e-05
propagation	str	The propagation method to use ('IBP' or 'SABR'). Default is 'IBP'.	'IBP'
sabr_args	dict	The arguments for the SABR method. Default is None.	None

Returns:

Name	Type	Description
taps_bound	Tuple[Tensor, Tensor]	The TAPS bounds of the model's output.

Source code in CTRAIN/bound/taps.py

def bound_taps(original_model, hardened_model, bounded_blocks, data, target, n_classes, ptb, device='cuda', pgd_steps=20, pgd_restarts=1, pgd_step_size=.2, 
               pgd_decay_factor=.2, pgd_decay_checkpoints=(5,7),
               gradient_link_thresh=.5, gradient_link_tolerance=1e-05, propagation="IBP", sabr_args=None):
    """
    Compute the bounds of the model's output using the TAPS method.

    Parameters:
        original_model (torch.nn.Module): The original neural network model.
        hardened_model (autoLiRPA.BoundedModule): The auto_LiRPA model.
        bounded_blocks (list): List of bounded blocks of the model.
        data (Tensor): The input data tensor.
        target (Tensor): The target labels tensor.
        n_classes (int): The number of classes for classification.
        ptb (auto_LiRPA.PerturbationLpNorm): The perturbation object defining the perturbation set.
        device (str, optional): The device to run the computation on. Default is 'cuda'.
        pgd_steps (int, optional): The number of steps for the PGD attack. Default is 20.
        pgd_restarts (int, optional): The number of restarts for the PGD attack. Default is 1.
        pgd_step_size (float, optional): The step size for the PGD attack. Default is 0.2.
        pgd_decay_factor (float, optional): The decay factor for the PGD attack. Default is 0.2.
        pgd_decay_checkpoints (tuple, optional): The decay checkpoints for the PGD attack. Default is (5, 7).
        gradient_link_thresh (float, optional): The threshold for gradient linking. Default is 0.5.
        gradient_link_tolerance (float, optional): The tolerance for gradient linking. Default is 1e-05.
        propagation (str, optional): The propagation method to use ('IBP' or 'SABR'). Default is 'IBP'.
        sabr_args (dict, optional): The arguments for the SABR method. Default is None.

    Returns:
        taps_bound(Tuple[Tensor, Tensor]): The TAPS bounds of the model's output.
    """
    assert len(bounded_blocks) == 2, "Split not supported!"

    if propagation == 'IBP':
        lb, ub = bound_ibp(
            model=bounded_blocks[0],
            ptb=ptb,
            data=data,
            target=None,
            n_classes=n_classes,
        )
    if propagation == 'SABR':
        assert sabr_args is not None, "Need to Provide SABR arguments if you choose SABR for propagation"
        lb, ub = bound_sabr(
            # Intermediate Bound model instructs to return bounds after the first network block
            **{**sabr_args, "intermediate_bound_model": bounded_blocks[0], "return_adv_output": False},
        )

    with torch.no_grad():
        hardened_model.eval()
        original_model.eval()
        for block in bounded_blocks:
            block.eval()
        c = construct_c(data, target, n_classes)
        with torch.no_grad():
            grad_cleaner = torch.optim.SGD(hardened_model.parameters())
            adv_samples = _get_pivotal_points(bounded_blocks[1], lb, ub, pgd_steps, pgd_restarts, pgd_step_size, pgd_decay_factor, pgd_decay_checkpoints, n_classes, C=c)
            grad_cleaner.zero_grad()

        hardened_model.train()
        original_model.train()
        for block in bounded_blocks:
            block.train()

    pts = adv_samples[0].detach()
    pts = torch.transpose(pts, 0, 1)
    pts = RectifiedLinearGradientLink.apply(lb.unsqueeze(0), ub.unsqueeze(0), pts, gradient_link_thresh, gradient_link_tolerance)
    pts = torch.transpose(pts, 0, 1)
    pgd_bounds = _get_bound_estimation_from_pts(bounded_blocks[1], pts, None, c)
    # NOTE: VERY IMPORTANT CHANGES TO TAPS BOUND TO BE COMPATIBLE WITH CTRAIN WORKFLOW
    pgd_bounds = pgd_bounds[:, 1:]
    pgd_bounds = -pgd_bounds


    ibp_lb, ibp_ub = bound_ibp(
        model=bounded_blocks[1],
        ptb=PerturbationLpNorm(x_L=lb, x_U=ub),
        data=data,
        target=target,
        n_classes=n_classes,
    )

    return pgd_bounds, ibp_lb