Model Functionalities

This section provides a detailed overview of the various modules within the core.model package of our project. Each module contains a set of functionalities that are integral to the operation of our model.

core.model.model

This module contains the main model functionalities.

This module defines a Variational Autoencoder (VAE) model with a unique architecture designed for image data. The model allows for the imposition of a sparsity constraint on certain layers, which can be useful for learning more compact or interpretable representations. The model is composed of an encoder, which uses a VGG19 architecture and includes a pixel shuffling operation, and a decoder, which gradually upscales the feature maps back to the original image size. The model also includes a method for accessing the activations of any layer.

Example usage:
    model = Model('model_name', 'cuda', [1, 2, 3], [1, 2], 1., 0.001)
    x = torch.randn(1, 3, 224, 224)
    output, kl, att = model(x)

class core.model.model.Model(model_name, device, layer_sparsity_cstraint=[], attention=[], sparsity_coeff=1.0, sparsity_param=0.001)

A Variational Autoencoder (VAE) model with a unique architecture designed for image data. The model allows for the imposition of a sparsity constraint on certain layers.

Parameters:

model_name (str) – The name of the model.
device (str) – The device to run the model on (‘cpu’ or ‘cuda’).
layer_sparsity_cstraint (list, optional) – A list of indices of layers to apply the sparsity constraint to.
attention (list, optional) – A list of attention values.
sparsity_coeff (float, optional) – The coefficient for the sparsity constraint.
sparsity_param (float, optional) – The parameter for the sparsity constraint.

__annotations__ = {}

__init__(model_name, device, layer_sparsity_cstraint=[], attention=[], sparsity_coeff=1.0, sparsity_param=0.001): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.model'

build_decoder()

build_encoder(layer_sparsity_cstraint, attention, sparsity_param, sparsity_coeff)

decode(x)

Decodes the input from the latent space.

Parameters:: x (torch.Tensor) – The input tensor.
Returns:: The decoded tensor.
Return type:: torch.Tensor

encode(x, sample=False)

Encodes the input into the latent space.

Parameters:

x (torch.Tensor) – The input tensor.
sample (bool, optional) – Whether to sample from the latent space.

Returns:

The encoded tensor.

Return type:

torch.Tensor

forward(x)

Passes the input through the model.

Parameters:: x (torch.Tensor) – The input tensor.
Returns:: A tuple containing the output tensor, the KL divergence, and the attention values.
Return type:: tuple

get_activations(x, layer_name)

Returns the activations of a specific layer.

Parameters:

x (torch.Tensor) – The input tensor.
layer_name (str) – The name of the layer.

Returns:

The activations of the layer.

Return type:

torch.Tensor

class core.model.model.PixelShuffler(upscale_factor)

A module that performs pixel shuffling, which is a way of upscaling the feature maps without introducing any new parameters. The upscale factor determines the factor by which the spatial dimensions are increased.

Parameters:: upscale_factor (int) – The factor by which to upscale the feature maps.

__annotations__ = {}

__init__(upscale_factor): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.model'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

core.model.basic_modules

This module contains the basic building blocks used in our model.

Modules used in building the main model

This module provides classes for building a neural network with optional sparsity constraints. It is designed to work with a pretrained VGG19 model, but can be adapted to other models as well.

Classes

NoConstraint: A module that applies a ReLU activation function to its input. (The encoder is built from ConvBlock, MeanStdFeatureMaps and Reparametrization layers. All of them have a add_constraint function that switches their NoConstraint to SparsityConstraint)
SparsityConstraint: A module that applies a sigmoid activation function to its input, then computes the KL divergence between the mean activation and a constant distribution. This is used to encourage sparsity in the activations.
ResBlock: A standard residual block, which consists of two convolutional layers with batch normalization and LeakyReLU activation, followed by an addition operation that adds the input to the output of the convolutions.
ConvBlock: A module that wraps a convolutional layer with a ReLU activation function and an optional sparsity constraint. The sparsity constraint can be added using the add_constraint method.
VGG19_Features: A module that extracts non-pretrained VGG19 feature extractor, organizes it in a module where you can apply the sparsity constraints anywhere as well as add attention layers.
MeanStdFeatureMaps: A module which generates a set of featuremaps for means and one for stds. It must be followed by a Reparametrization module.
Reparametrization: A module which performs re-parametrization (sampling from a gaussian, then x sig + mu). The module right before must be a MeanStdFeatureMaps() module.

class core.model.basic_modules.ConvBlock(conv2d, idx)

A PyTorch module that wraps a convolutional layer with a ReLU activation function and an optional sparsity constraint. The sparsity constraint can be added using the add_constraint method.

Attributes

name: str: The name of the ConvBlock, which includes the index and the activation function used.
conv: torch.nn.Conv2d: The convolutional layer wrapped by the ConvBlock.
constraint: torch.nn.Module: The constraint applied to the output of the convolutional layer. This is initially a NoConstraint instance, but can be replaced with a SparsityConstraint instance using the add_constraint method.

Methods

forward(x: torch.Tensor) -> Tuple[torch.Tensor, float]: Applies the convolutional layer and the constraint to the input tensor and returns the result.
add_constraint(constraint_param: float, scaling_param: float) -> None: Replaces the current constraint with a SparsityConstraint instance using the provided parameters.
get_activations(x: torch.Tensor) -> torch.Tensor: Applies the convolutional layer and the constraint to the input tensor and returns the result without gradients.

__annotations__ = {}

__init__(conv2d, idx): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.basic_modules'

add_constraint(constraint_param, scaling_param)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_activations(x)

class core.model.basic_modules.MeanStdFeatureMaps(in_channels, latent_dim)

A PyTorch module which generates a set of feature maps for means and one for standard deviations. It must be followed by a Reparametrization module.

Attributes

name: str: The name of the MeanStdFeatureMaps module, which includes the activation function used.
conv_means: torch.nn.Conv2d: The convolutional layer that generates the mean feature maps.
conv_stds: torch.nn.Conv2d: The convolutional layer that generates the standard deviation feature maps.
constraint: torch.nn.Module: The constraint applied to the output of the convolutional layers. This is initially a NoConstraint instance, but can be replaced with a SparsityConstraint instance using the add_constraint method.

Methods

forward(x: torch.Tensor) -> Tuple[torch.Tensor, float]: Applies the convolutional layers and the constraint to the input tensor and returns the result.
add_constraint(constraint_param: float, scaling_param: float) -> None: Replaces the current constraint with a SparsityConstraint instance using the provided parameters.
get_activations(x: torch.Tensor, mean_or_std: str) -> torch.Tensor: Applies the appropriate convolutional layer and the constraint to the input tensor and returns the result without gradients.

__annotations__ = {}

__init__(in_channels, latent_dim): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.basic_modules'

add_constraint(constraint_param, scaling_param)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_activations(x, mean_or_std)

class core.model.basic_modules.NoConstraint

A PyTorch module that applies a ReLU activation function to its input. This class is used as a placeholder when no sparsity constraint is desired.

Attributes

None

Methods

forward(x: torch.Tensor) -> Tuple[torch.Tensor, float]: Applies the ReLU activation function to the input tensor and returns the result along with 0.0.

__annotations__ = {}

__init__(): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.basic_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class core.model.basic_modules.Reparametrization(device)

A PyTorch module which performs re-parametrization (sampling from a Gaussian, then x sig + mu). The module right before must be a MeanStdFeatureMaps() module.

Attributes

name: str: The name of the Reparametrization module, which includes the activation function used.
constraint: torch.nn.Module: The constraint applied to the output of the re-parametrization. This is initially a NoConstraint instance, but can be replaced with a SparsityConstraint instance using the add_constraint method.
normal: torch.distributions.Normal: A Normal distribution used for sampling in the re-parametrization.

Methods

forward(x: torch.Tensor) -> Tuple[torch.Tensor, float]: Performs the re-parametrization on the input tensor and returns the result.
add_constraint(constraint_param: float, scaling_param: float) -> None: Replaces the current constraint with a SparsityConstraint instance using the provided parameters.
get_activations(x: torch.Tensor) -> torch.Tensor: Performs the re-parametrization on the input tensor and returns the result without gradients.

__annotations__ = {}

__init__(device): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.basic_modules'

add_constraint(constraint_param, scaling_param)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_activations(x)

class core.model.basic_modules.ResBlock(nb_chan)

A standard residual block, which consists of two convolutional layers with batch normalization and LeakyReLU activation, followed by an addition operation that adds the input to the output of the convolutions.

Attributes

operations: torch.nn.Sequential: A sequence of operations that define the residual block.

Methods

forward(x: torch.Tensor) -> torch.Tensor: Applies the operations of the residual block to the input tensor and returns the result.

__annotations__ = {}

__init__(nb_chan): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.basic_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class core.model.basic_modules.SparsityConstraint(constraint_param, scaling_param)

A PyTorch module that applies a sigmoid activation function to its input, then computes the KL divergence between the mean activation and a constant distribution. This is used to encourage sparsity in the activations.

Attributes

constraint_param: float: The parameter for the constant distribution in the KL divergence calculation.
scaling_param: float: The scaling factor for the KL divergence.

Methods

forward(x: torch.Tensor) -> Tuple[torch.Tensor, float]: Applies the sigmoid activation function to the input tensor, computes the KL divergence, and returns both the activated tensor and the scaled KL divergence.

__annotations__ = {}

__init__(constraint_param, scaling_param): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.basic_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class core.model.basic_modules.VGG19_Features(constraint_locations, attention, constraint_param, scaling_param)

A PyTorch module that extracts a non-pretrained VGG19 feature extractor, organizes it in a module where you can apply the sparsity constraints anywhere as well as add attention layers.

Attributes

name: str: The name of the VGG19_Features module, which includes the activation function used.
feature_extractor: torch.nn.Sequential: A sequence of operations that define the feature extraction process.
layers: List[str]: A list of names of the ConvBlock layers in the feature extractor.
attention_layers: List[str]: A list of names of the Constrained_CBAM layers in the feature extractor.
all_layers: List[str]: A list of names of all the ConvBlock and Constrained_CBAM layers in the feature extractor.

Methods

forward(x: torch.Tensor) -> Tuple[torch.Tensor, float, float]: Applies the feature extraction process to the input tensor and returns the result, along with the total KL divergence and attention loss.
get_activations(x: torch.Tensor, layer_name: str) -> torch.Tensor: Applies the feature extraction process up to the specified layer to the input tensor and returns the result.
__len__() -> int: Returns the number of layers in the feature extractor.

__annotations__ = {}

__init__(constraint_locations, attention, constraint_param, scaling_param): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__len__()

__module__ = 'core.model.basic_modules'

add_attention_layers(attention, constraint_param, scaling_param)

build_convs(features, constraint_locations, constraint_param, scaling_param)

build_feature_extractor(features, constraint_locations, constraint_param, scaling_param)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_activations(x, layer_name)

core.model.attention_modules

This module contains the attention mechanisms used in our model.

class core.model.attention_modules.BasicConv(in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False)

This class is a basic convolutional layer, taken from git.. It includes options for batch normalization and ReLU activation.

__annotations__ = {}

__init__(in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.attention_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class core.model.attention_modules.ChannelGate(gate_channels, reduction_ratio=16, pool_types=['avg', 'max'])

This class implements a channel gate, taken from git.. It includes options for different types of pooling.

__annotations__ = {}

__init__(gate_channels, reduction_ratio=16, pool_types=['avg', 'max']): Initializes internal Module state, shared by both nn.Module and ScriptModule.

__module__ = 'core.model.attention_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_att_map(x)

class core.model.attention_modules.ChannelPool(*args, **kwargs)

This class implements a channel pooling layer, taken from git.. It returns the concatenation of the max and mean of its input.

__annotations__ = {}

__module__ = 'core.model.attention_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class core.model.attention_modules.Constrained_CBAM(gate_channels, constraint_param, scaling_param, index, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False)

This class implements a modified version of the Convolutional Block Attention Module (CBAM) with a L1 constraint. It only includes depth-wise attention, also known as channel attention.

Attributes:: gate_channels (int): The number of output channels in the attention module. constraint_param (float): The parameter for the L1 constraint. scaling_param (float): The scaling parameter for the attention module. index (int): The index of the attention module in the model. reduction_ratio (int, optional): The ratio for the channel reduction in the attention module. Default is 16. pool_types (list of str, optional): The types of pooling to use in the attention module. Default is [‘avg’, ‘max’]. no_spatial (bool, optional): Whether to exclude spatial attention. Default is False.

Example usage:

from model import Constrained_CBAM

# Initialize the Constrained_CBAM module
cbam = Constrained_CBAM(gate_channels=64, constraint_param=0.1, scaling_param=0.5, index=1)

# Use the Constrained_CBAM module on an input
output, l1 = cbam(input)

__annotations__ = {}

__init__(gate_channels, constraint_param, scaling_param, index, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False)

Initialize the Constrained_CBAM module.

Parameters:: gate_channels (int): The number of output channels in the attention module. constraint_param (float): The parameter for the L1 constraint. scaling_param (float): The scaling parameter for the attention module. index (int): The index of the attention module in the model. reduction_ratio (int, optional): The ratio for the channel reduction in the attention module. Default is 16. pool_types (list of str, optional): The types of pooling to use in the attention module. Default is [‘avg’, ‘max’]. no_spatial (bool, optional): Whether to exclude spatial attention. Default is False.

__module__ = 'core.model.attention_modules'

forward(x)

Forward pass of the Constrained_CBAM module.

Parameters:: x (torch.Tensor): The input tensor.
Returns:: tuple: The output tensor and the L1 norm.

get_activations(x)

Get the activation map of the Constrained_CBAM module.

Parameters:: x (torch.Tensor): The input tensor.
Returns:: torch.Tensor: The activation map.

class core.model.attention_modules.Flatten(*args, **kwargs)

This class is a simple layer that flattens its input, taken from git..

__annotations__ = {}

__module__ = 'core.model.attention_modules'

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

core.model.attention_modules.logsumexp_2d(tensor): This function calculates the log-sum-exp of a 2D tensor, taken from git.