4. Functionality for reproducibility

The following page describe the modules available on the T-SHIELD platform.

TSHIELD.TSHIELD.tshield
TSHIELD.procedures.procedures.classifier(...)	Model used for classification on the T-SHIELD work [SGarciaLH24] .
TSHIELD.procedures.procedures.train_step(...)	Train step for the model.
TSHIELD.procedures.procedures.validation_step(...)	Validation step for the model.

TSHIELD.TSHIELD.rshield(model: Module, input: Tensor, input_0: Tensor = None, segmentation: int = 1, device: device | str = device(type='cpu'), percentage=None)

View Source Code

def rshield(
    model: torch.nn.Module,
    input: torch.Tensor,
    input_0: torch.Tensor = None,
    segmentation: int = 1,
    device: Union[torch.device, str] = torch.device("cpu"),
    percentage=None,
):
    """
    Random-Selective Hidden Input Evaluation for Learning Dynamics (R-SHIELD).
    
    This function calculates the R-SHIELD score for a given model and input tensor.
    R-SHIELD is a method for evaluating the importance of different regions in an input tensor
    for the predictions made by a model. It measures the sensitivity of the model's output to changes
    in specific regions of the input. Those regions are randomized so the method is stochastic.

    :param model: The model to be evaluated.
    :param input: The input tensor for which the SHIELD score is calculated.
    :param input_0: The input tensor to be considered as the `null` input with the same shape as input.
    :param segmentation: The segmentation method to be used.
    :param device: The device to be used for computation.
    :param percentage: The percentage of the input to be evaluated.
    :return: The `r-shield_score`. It is a scalar tensor representing the R-SHIELD score for the given input tensor and can
    be used as a regularization term in the training of the model given as parameter.

    Example usage:

    .. code-block:: python

        from TSHIELD.TSHIELD import rshield
        model = MyModel() \\ Your classification model
        input = torch.rand((1, 3, 224, 224))
        rshield_score = rshield(model, input, input_0=None, segmentation=1, device='cuda', percentage=2)
        
    """

    if input_0 == None:
        input_0 = torch.zeros_like(input) + torch.mean(
            input, dim=(1, 2, 3), keepdim=True
        )
    model = model.to(device)
    input = input.to(device)
    input_0 = input_0.to(device)

    output_original = checkpoint(model, input,use_reentrant=True)

    feat_importance = (
        torch.rand((input.shape[0], input.shape[2], input.shape[3])) * 2 - 1
    )
    feat_importance = feat_importance.to(device)

    feat_importance_max = torch.nn.MaxPool2d(
        kernel_size=segmentation, stride=segmentation
    )(feat_importance)
    feat_importance_min = torch.nn.MaxPool2d(
        kernel_size=segmentation, stride=segmentation
    )(-feat_importance)

    feat_importance_final = torch.where(
        feat_importance_max > -feat_importance_min,
        feat_importance_max,
        -feat_importance_min,
    )
    feat_importance_final = feat_importance_final.unsqueeze(1)

    feat_importance_final = torch.nn.UpsamplingNearest2d(size=input.shape[2:])(
        feat_importance_final
    )
    feat_importance_final = feat_importance_final.squeeze(1)
    abs_importance = torch.abs(feat_importance_final)
    quantile_abs = torch.quantile(
        abs_importance, q=percentage / 100.0, dim=1, keepdim=True
    )
    quantile_abs = torch.quantile(
        quantile_abs, q=percentage / 100.0, dim=2, keepdim=True
    )

    quantile_abs = quantile_abs.repeat(
        1, abs_importance.shape[1], abs_importance.shape[2]
    )

    mask = torch.where(
        quantile_abs > abs_importance,
        torch.ones_like(quantile_abs),
        torch.zeros_like(quantile_abs),
    )
    mask = mask.unsqueeze(1)
    mask = mask.repeat(1, input.shape[1], 1, 1)

    modified_input = torch.where(mask == 0, input, input_0)

    modified_output = checkpoint(model, modified_input,use_reentrant=True)

    Px_modif = torch.softmax(modified_output, dim=1)
    Px = torch.softmax(output_original, dim=1)

    constraint_1 = torch.mean(
        -torch.sum(Px * (torch.log(Px_modif) - torch.log(Px)), dim=1)
    )
    constraint_2 = torch.mean(
        -torch.sum(Px_modif * (torch.log(Px) - torch.log(Px_modif)), dim=1)
    )

    constraint = constraint_1 + constraint_2

    if torch.isnan(constraint):
        constraint = torch.tensor(0.0).to(device)

    return constraint

Random-Selective Hidden Input Evaluation for Learning Dynamics (R-SHIELD).

This function calculates the R-SHIELD score for a given model and input tensor. R-SHIELD is a method for evaluating the importance of different regions in an input tensor for the predictions made by a model. It measures the sensitivity of the model’s output to changes in specific regions of the input. Those regions are randomized so the method is stochastic.

Parameters:

• model – The model to be evaluated.

• input – The input tensor for which the SHIELD score is calculated.

• input_0 – The input tensor to be considered as the null input with the same shape as input.

• segmentation – The segmentation method to be used.

• device – The device to be used for computation.

• percentage – The percentage of the input to be evaluated.

Returns:

The r-shield_score. It is a scalar tensor representing the R-SHIELD score for the given input tensor and can

be used as a regularization term in the training of the model given as parameter.

Example usage:

from TSHIELD.TSHIELD import rshield
model = MyModel() \ Your classification model
input = torch.rand((1, 3, 224, 224))
rshield_score = rshield(model, input, input_0=None, segmentation=1, device='cuda', percentage=2)

TSHIELD.TSHIELD.xshield(model: Module, input: Tensor, input_0: Tensor = None, segmentation: int = 1, device: device | str = device(type='cpu'), percentage=None)

View Source Code

def xshield(
    model: torch.nn.Module,
    input: torch.Tensor,
    input_0: torch.Tensor = None,
    segmentation: int = 1,
    device: Union[torch.device, str] = torch.device("cpu"),
    percentage=None,
):
    """
    xAI - Selective Hidden Input Evaluation for Learning Dynamics (X-SHIELD).

    This function calculates the X-SHIELD score for a given model and input tensor.
    X-SHIELD is a method based on previous SHIELD method for evaluating the importance of different 
    regions in an input tensor for the predictions made by a model. It measures the sensitivity of 
    the model's output to changes in specific regions of the input.
    
    X-SHIELD uses an explanation defined by the model itself to evaluate the importance of different regions in the input tensor and ocludes the regions with less importance.

    :param model: The model to be evaluated.
    :param input: The input tensor for which the X-SHIELD score is calculated.
    :param input_0: The input tensor to be considered as the `null` input with the same shape as input.
    :param segmentation: The segmentation method to be used.
    :param device: The device to be used for computation.
    :param percentage: The percentage of the input to be evaluated.
    :return: The `xshield_score`. It is a scalar tensor representing the X-SHIELD score for the given input tensor and can
    be used as a regularization term in the training of the model given as parameter.

    Example usage:

    .. code-block:: python

        from TSHIELD.TSHIELD import xshield
        model = MyModel() \\ Your classification model
        input = torch.rand((1, 3, 224, 224))
        xshield_score = xshield(model, input, input_0=None, segmentation=1, device='cuda', percentage=2)

    """

    if input_0 == None:
        input_0 = torch.zeros_like(input) + torch.mean(
            input, dim=(1, 2, 3), keepdim=True
        )
    model = model.to(device)
    input = input.to(device)
    input_0 = input_0.to(device)

    output_original = checkpoint(model, input, use_reentrant=True)

    feat_importance = feature_importance(model,input,forward_input=output_original,device=device)
    feat_importance = feat_importance.to(device)

    feat_importance_max = torch.nn.MaxPool2d(
        kernel_size=segmentation, stride=segmentation
    )(feat_importance)
    feat_importance_min = torch.nn.MaxPool2d(
        kernel_size=segmentation, stride=segmentation
    )(-feat_importance)

    feat_importance_final = torch.where(
        feat_importance_max > -feat_importance_min,
        feat_importance_max,
        -feat_importance_min,
    )
    feat_importance_final = feat_importance_final.unsqueeze(1)

    feat_importance_final = torch.nn.UpsamplingNearest2d(size=input.shape[2:])(
        feat_importance_final
    )
    feat_importance_final = feat_importance_final.squeeze(1)
    abs_importance = torch.abs(feat_importance_final)
    quantile_abs = torch.quantile(
        abs_importance, q=percentage / 100.0, dim=1, keepdim=True
    )
    quantile_abs = torch.quantile(
        quantile_abs, q=percentage / 100.0, dim=2, keepdim=True
    )

    quantile_abs = quantile_abs.repeat(
        1, abs_importance.shape[1], abs_importance.shape[2]
    )

    mask = torch.where(
        quantile_abs > abs_importance,
        torch.ones_like(quantile_abs),
        torch.zeros_like(quantile_abs),
    )
    mask = mask.unsqueeze(1)
    mask = mask.repeat(1, input.shape[1], 1, 1)

    modified_input = torch.where(mask == 0, input, input_0)

    modified_output = checkpoint(model, modified_input, use_reentrant=True)

    Px_modif = torch.softmax(modified_output, dim=1)
    Px = torch.softmax(output_original, dim=1)

    constraint_1 = torch.mean(
        -torch.sum(Px * (torch.log(Px_modif) - torch.log(Px)), dim=1)
    )
    constraint_2 = torch.mean(
        -torch.sum(Px_modif * (torch.log(Px) - torch.log(Px_modif)), dim=1)
    )

    constraint = constraint_1 + constraint_2

    if torch.isnan(constraint):
        constraint = torch.tensor(0.0).to(device)

    return constraint

xAI - Selective Hidden Input Evaluation for Learning Dynamics (X-SHIELD).

This function calculates the X-SHIELD score for a given model and input tensor. X-SHIELD is a method based on previous SHIELD method for evaluating the importance of different regions in an input tensor for the predictions made by a model. It measures the sensitivity of the model’s output to changes in specific regions of the input.

X-SHIELD uses an explanation defined by the model itself to evaluate the importance of different regions in the input tensor and ocludes the regions with less importance.

Parameters:

• model – The model to be evaluated.

• input – The input tensor for which the X-SHIELD score is calculated.

• input_0 – The input tensor to be considered as the null input with the same shape as input.

• segmentation – The segmentation method to be used.

• device – The device to be used for computation.

• percentage – The percentage of the input to be evaluated.

Returns:

The xshield_score. It is a scalar tensor representing the X-SHIELD score for the given input tensor and can

be used as a regularization term in the training of the model given as parameter.

Example usage:

from TSHIELD.TSHIELD import xshield
model = MyModel() \ Your classification model
input = torch.rand((1, 3, 224, 224))
xshield_score = xshield(model, input, input_0=None, segmentation=1, device='cuda', percentage=2)

TSHIELD.procedures.procedures.classifier(pretrained_model, num_classes)

View Source Code

def classifier(pretrained_model, num_classes):

    # cite sevillano2024shield:
    """
    Model used for classification on the T-SHIELD work :cite:`sevillano2024shield` .
    This function returns a model with the specified number of classes and the specified pretrained model of
    the torchvision library.

    This method was used as a wrapper to load the classifier but
    you can use your own model if it accepts the same input format as the torchvision
    models and returns a tensor with the same shape as the number of classes in the logits space.

    :param pretrained_model: The name of the pretrained model to be used. It can be one of the following:
        - efficientnet-b2
        - efficientnet_v2_s
        - vit_b_16
        - swin_v2_s
    :param num_classes: The number of classes of the dataset.
    :return: The model with the specified number of classes and the specified pretrained model. We
        delete the last layer of the model and add a new layer with the specified number of classes.

    To load a specific model, use the following code:

    .. code-block:: python

        model = classifier("efficientnet-b2", 10)
    """

    if "efficientnet-b2" == pretrained_model:
        model_pretrained = torchvision.models.efficientnet_b2(weights="IMAGENET1K_V1")
        pretrained_state_dict = model_pretrained.state_dict()

        model = torchvision.models.efficientnet_b2(num_classes=num_classes)
        state_dict = model.state_dict()

        pretrained_state_dict["classifier.1.weight"] = state_dict["classifier.1.weight"]
        pretrained_state_dict["classifier.1.bias"] = state_dict["classifier.1.bias"]
        model.load_state_dict(pretrained_state_dict)
    elif "efficientnet_v2_s" == pretrained_model:
        model_pretrained = torchvision.models.efficientnet_v2_s(weights="IMAGENET1K_V1")
        pretrained_state_dict = model_pretrained.state_dict()

        model = torchvision.models.efficientnet_v2_s(num_classes=num_classes)
        state_dict = model.state_dict()

        pretrained_state_dict["classifier.1.weight"] = state_dict["classifier.1.weight"]
        pretrained_state_dict["classifier.1.bias"] = state_dict["classifier.1.bias"]
        model.load_state_dict(pretrained_state_dict)
    elif "vit_b_16" == pretrained_model:
        model_pretrained = torchvision.models.vit_b_16(weights="IMAGENET1K_V1")
        pretrained_state_dict = model_pretrained.state_dict()

        model = torchvision.models.vit_b_16(num_classes=num_classes)
        state_dict = model.state_dict()

        pretrained_state_dict["heads.head.weight"] = state_dict["heads.head.weight"]
        pretrained_state_dict["heads.head.bias"] = state_dict["heads.head.bias"]
        model.load_state_dict(pretrained_state_dict)

    elif "swin_v2_s" == pretrained_model:
        model_pretrained = torchvision.models.swin_v2_s(weights="IMAGENET1K_V1")
        pretrained_state_dict = model_pretrained.state_dict()
        model = torchvision.models.swin_v2_s(num_classes=num_classes)
        state_dict = model.state_dict()

        pretrained_state_dict["head.weight"] = state_dict["head.weight"]
        pretrained_state_dict["head.bias"] = state_dict["head.bias"]
        model.load_state_dict(pretrained_state_dict)
    return model

Model used for classification on the T-SHIELD work [SGarciaLH24] . This function returns a model with the specified number of classes and the specified pretrained model of the torchvision library.

This method was used as a wrapper to load the classifier but you can use your own model if it accepts the same input format as the torchvision models and returns a tensor with the same shape as the number of classes in the logits space.

Parameters:

• pretrained_model – The name of the pretrained model to be used. It can be one of the following: - efficientnet-b2 - efficientnet_v2_s - vit_b_16 - swin_v2_s

• num_classes – The number of classes of the dataset.

Returns:

The model with the specified number of classes and the specified pretrained model. We delete the last layer of the model and add a new layer with the specified number of classes.

To load a specific model, use the following code:

model = classifier("efficientnet-b2", 10)

TSHIELD.procedures.procedures.train_step(ds_loader, model, optimizer, loss_f, reg_f, device, transform=None, train=True)

View Source Code

def train_step(
    ds_loader, model, optimizer, loss_f, reg_f, device, transform=None, train=True
):
    '''
    Train step for the model. This function proceed to do a full epoch of training or validation, detending
    on the `train` parameter. It returns the loss, accuracy and regularization of the model.
    
    :param ds_loader: The dataloader to be used for training or validation.
    :param model: The model to be trained or validated.
    :param optimizer: The optimizer to be used for training. If `train` is False, this parameter is not used.
    :param loss_f: The loss function to be used for training or validation.
    :param reg_f: The regularization function to be used for training or validation. It can be None for no regularization.
    :param device: The device to be used for training or validation.
    :param transform: The transformation to be used for training. If `train` is False, this parameter is not used. It 
        can be None for no data augmentation or transformation needed.
    :param train: A boolean indicating if the model should be trained or validated. 
        - If `train` is True, the model is trained and updated with the optimizer.
        - If `train` is False, the model is validated and the optimizer is not used.
    :return: The loss, accuracy and regularization of the model.
    
    Example of a full step of training and validation using this function:
    
    .. code-block:: python
    
        # We have a model, a train dataloader, a validation dataloader, a loss function and a regularization function. 
        # The previous step had `best_val_loss` as a variable to store the best validation loss over the epochs.
        
        # The model weights are updated with the optimizer on the training phase. The gradient is calculated with the 
        # loss and the regularization function.
        loss, acc, reg = train_step(train_dataloader, model, optimizer, loss_f, reg_f, device, transform, train=True)
        
        # The model is not updated on the validation phase. The best model is saved if the validation `loss+reg` is 
        # better than the previous one.
        val_loss, val_acc, val_reg = train_step(validation_dataloader, model, optimizer, loss_f, reg_f, device, transform, train=False)
        if loss+reg < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), "best_model.pt")
    '''
    ACC, LOSS, REGS = 0.0, 0.0, 0.0

    ds_loader = tqdm.tqdm(ds_loader, desc="Training" if train else "Validation")
    total = 0
    if train == True:
        model.train()
    else:
        model.eval()
    torch.cuda.empty_cache()
    model = model.to(device)
    for data in ds_loader:
        batch_input, batch_labels = data
        total += len(batch_input)
        batch_input, batch_labels = batch_input.float().to(device), batch_labels.to(device)
        if transform != None and train == True:
            batch_input = transform(batch_input)
        batch_input.requires_grad = True
        batch_input, batch_labels = batch_input.float().to(device), batch_labels.to(device)

        if train == True:
            optimizer.zero_grad()
        reg = reg_f(model, batch_input)
        # output = model(batch_input)
        output = checkpoint(model, batch_input,use_reentrant=True)

        loss = loss_f(output, batch_labels)
        LOSS += loss.item()

        loss += reg if reg != None else 0
        if train == True:
            loss.backward()
            optimizer.step()

        ACC += torch.sum(
            (torch.argmax(output, dim=-1) == torch.argmax(batch_labels, dim=-1)).float()
        ).item()
        if isinstance(reg, torch.Tensor):
            REGS += reg.item()
        else:
            REGS +=  reg if reg != None else 0

        ds_loader.set_postfix(
            {
                "loss": LOSS / total,
                "acc": ACC / total,
                "reg": REGS / total,
            }
        )

    return (
        LOSS / total,
        ACC / total,
        REGS / total,
    )

Train step for the model. This function proceed to do a full epoch of training or validation, detending on the train parameter. It returns the loss, accuracy and regularization of the model.

Parameters:

• ds_loader – The dataloader to be used for training or validation.

• model – The model to be trained or validated.

• optimizer – The optimizer to be used for training. If train is False, this parameter is not used.

• loss_f – The loss function to be used for training or validation.

• reg_f – The regularization function to be used for training or validation. It can be None for no regularization.

• device – The device to be used for training or validation.

• transform – The transformation to be used for training. If train is False, this parameter is not used. It can be None for no data augmentation or transformation needed.

• train – A boolean indicating if the model should be trained or validated. - If train is True, the model is trained and updated with the optimizer. - If train is False, the model is validated and the optimizer is not used.

Returns:

The loss, accuracy and regularization of the model.

Example of a full step of training and validation using this function:

# We have a model, a train dataloader, a validation dataloader, a loss function and a regularization function.
# The previous step had `best_val_loss` as a variable to store the best validation loss over the epochs.

# The model weights are updated with the optimizer on the training phase. The gradient is calculated with the
# loss and the regularization function.
loss, acc, reg = train_step(train_dataloader, model, optimizer, loss_f, reg_f, device, transform, train=True)

# The model is not updated on the validation phase. The best model is saved if the validation `loss+reg` is
# better than the previous one.
val_loss, val_acc, val_reg = train_step(validation_dataloader, model, optimizer, loss_f, reg_f, device, transform, train=False)
if loss+reg < best_val_loss:
    best_val_loss = val_loss
    torch.save(model.state_dict(), "best_model.pt")

TSHIELD.procedures.procedures.validation_step(ds_loader, model, loss_f, reg_f, device)

View Source Code

def validation_step(ds_loader, model, loss_f, reg_f, device):
    '''
    Validation step for the model. 
    
    This function is a wrapper for the `train_step` function with the `train` parameter set to False.
    
    :param ds_loader: The dataloader to be used for validation.
    :param model: The model to be validated.
    :param loss_f: The loss function to be used for validation.
    :param reg_f: The regularization function to be used for validation. It can be None for no regularization.
    :param device: The device to be used for validation.
    :return: The loss, accuracy and regularization of the model.
    
    '''
    
    return train_step(
        ds_loader=ds_loader,
        model=model,
        optimizer=None,
        loss_f=loss_f,
        reg_f=reg_f,
        device=device,
        train=False,
    )

Validation step for the model.

This function is a wrapper for the train_step function with the train parameter set to False.

Parameters:

• ds_loader – The dataloader to be used for validation.

• model – The model to be validated.

• loss_f – The loss function to be used for validation.

• reg_f – The regularization function to be used for validation. It can be None for no regularization.

• device – The device to be used for validation.

Returns:

The loss, accuracy and regularization of the model.