Metrics and Utilities

This section provides a detailed overview of the various modules within the core folder which contain a set of functionalities related to data handling, metrics calculation, plotting, and model training.

core.data

This module contains functionalities related to data handling and preprocessing.

class core.data.OptionalSplitDataset(root_dir, split='train', device=None)

A custom PyTorch Dataset for loading image data.

This class is used to create a dataset object that can be used with a PyTorch DataLoader for efficient data loading.

Attributes:

device (torch.device): The device to load the images onto. file_paths (list): A list of paths to the image files.

Args:

root_dir (str): The root directory containing the image files. split (str): The dataset split (‘train’: keep 80% or ‘test’ (keep the other 20%) or ‘none’ (keep 100%)). device (torch.device): The device to load the images onto.

# Example 1: Creating a dataset for training images on the CPU
train_dataset = OptionalSplitDataset('/path/to/images', split='train', device=torch.device('cpu'))

# Example 2: Creating a dataset for test images on the GPU
test_dataset = OptionalSplitDataset('/path/to/images', split='test', device=torch.device('cuda'))

# Example 3: Accessing an image from the dataset
image = train_dataset[0]

# Example 4: Getting the total number of images in the dataset
num_images = len(train_dataset)

# Example 5: Creating a dataset without splitting the data
full_dataset = OptionalSplitDataset('/path/to/images', split='none', device=torch.device('cpu'))
__getitem__(index)

Get an image by index.

This method is used to get an image and its label by index.

Args:

index (int): The index of the image.

Returns:

torch.Tensor: The image tensor.

__init__(root_dir, split='train', device=None)

The constructor for OptionalSplitDataset class.

This method initializes the dataset object, it sets the device where the images will be loaded, gets all the image file paths from the root directory and sets the split for the dataset.

Args:

root_dir (str): The root directory containing the image files. split (str): The dataset split (‘train’ or ‘test’). device (torch.device): The device to load the images onto.

__len__()

Get the total number of images in the dataset.

This method is used to get the total number of images in the dataset.

Returns:

int: The total number of images.

__module__ = 'core.data'
__parameters__ = ()
get_all_image_paths(root_dir)

Recursively get all image paths from the root directory.

This method is used to get all the image file paths from the root directory.

Args:

root_dir (str): The root directory containing the image files.

Returns:

list: A list of paths to the image files.

static process_image(image_path, device)

Process an image.

This method reads an image from a file, checks its size, resizes it to 224x224, adds a batch dimension, and moves it to the specified device.

Args:

image_path (str): The path to the image file. device (torch.device): The device to move the image to.

Returns:

torch.Tensor: The processed image tensor.

core.metrics

This module contains all the metrics-extraction code.

This module contains the MetricExtractor class, a utility for extracting various metrics from a given model.

class core.metrics.MetricExtractor(model)

A class for extracting various metrics from a given model.

This class is initialized with a model and sets up several loss functions. It is primarily used to calculate and return a dictionary of metrics for a given image. These metrics include layer-wise activations, L1 norms, Gini coefficients, kurtosis, and reconstruction errors based on different loss functions. The class also calculates the Stochastic Weight Averaging Gaussian (SWAG) delta for each loss function.

Attributes:

model (nn.Module): The PyTorch model from which to extract metrics. loss_functions (dict): A dictionary of loss functions.

Example usage:

from core.model.model import Model

# Initialize a model
model = Model(model_name='my_model', device='cuda')

# Initialize the MetricExtractor
metric_extractor = MetricExtractor(model)

# Use the MetricExtractor on an image
metrics = metric_extractor('path_to_image.jpg')
__call__(image_path)

Calculate and return a dictionary of metrics for a given image.

Parameters:

image_path (str): The path to the image.

Returns:

dict: A dictionary of metrics.

__dict__ = mappingproxy({'__module__': 'core.metrics', '__doc__': "\n    A class for extracting various metrics from a given model.\n\n    This class is initialized with a model and sets up several loss functions. It is primarily used to calculate and return a dictionary of metrics for a given image. These metrics include layer-wise activations, L1 norms, Gini coefficients, kurtosis, and reconstruction errors based on different loss functions. The class also calculates the Stochastic Weight Averaging Gaussian (SWAG) delta for each loss function.\n\n    **Attributes**:\n        model (nn.Module): The PyTorch model from which to extract metrics.\n        loss_functions (dict): A dictionary of loss functions.\n\n    **Example usage**:\n\n    .. code-block:: python\n\n        from core.model.model import Model\n\n        # Initialize a model\n        model = Model(model_name='my_model', device='cuda')\n\n        # Initialize the MetricExtractor\n        metric_extractor = MetricExtractor(model)\n\n        # Use the MetricExtractor on an image\n        metrics = metric_extractor('path_to_image.jpg')\n\n    ", '__init__': <function MetricExtractor.__init__>, '__call__': <function MetricExtractor.__call__>, 'get_activations': <function MetricExtractor.get_activations>, 'calculate_metrics': <function MetricExtractor.calculate_metrics>, 'calculate_gini': <function MetricExtractor.calculate_gini>, 'get_reco_error': <function MetricExtractor.get_reco_error>, 'get_SAM_delta': <function MetricExtractor.get_SAM_delta>, '__dict__': <attribute '__dict__' of 'MetricExtractor' objects>, '__weakref__': <attribute '__weakref__' of 'MetricExtractor' objects>, '__annotations__': {}})
__init__(model)

Initialize the MetricExtractor with a model.

Parameters:

model (nn.Module): The PyTorch model from which to extract metrics.

__module__ = 'core.metrics'
__weakref__

list of weak references to the object (if defined)

calculate_gini(activations)

Calculate the reconstruction error for the current image using a specified loss function.

Parameters:

loss_name (str): The name of the loss function to use.

Returns:

dict: A dictionary with the loss name as the key and the reconstruction error as the value.

calculate_metrics(activations, layer)

Calculate and return a dictionary of metrics for a given layer’s activations.

The metrics calculated depend on whether the layer is an attention layer or not. For attention layers, the sum of absolute activations is calculated. For other layers, the L1 norm, Gini coefficient, and kurtosis of the activations are calculated.

Parameters:

activations (torch.Tensor): The activations of the layer. layer (str): The name of the layer.

Returns:

dict: A dictionary of metrics for the layer.

get_SAM_delta(loss_name)

Calculate the Sharpness-aware-Minimization -inspired delta for a specified loss function.

This metric measures how much a small change in model parameters can collapse its ability to reconstruct an image’

Parameters:

loss_name (str): The name of the loss function to use.

Returns:

dict: A dictionary with the loss name as the key and the sam delta as the value.

get_activations(layer)

Get the activations of a given layer for the current image.

Parameters:

layer (str): The name of the layer.

Returns:

torch.Tensor: The activations of the layer.

get_reco_error(loss_name)

Calculate the reconstruction error for the current image using a specified loss function.

Parameters:

loss_name (str): The name of the loss function to use.

Returns:

dict: A dictionary with the loss name as the key and the reconstruction error as the value.

core.plotting

This module contains functionalities for visualizing images and model histograms.

core.plotting.plot_activation_histograms(metric_extractor, dataset, file_path, num_images=5)

Plot activation histograms for a given number of images from a dataset.

This function takes a MetricExtractor instance, a dataset, a file path, and an optional number of images to plot. It then plots activation histograms for each layer of the model for each image, and saves the plot to a specified file.

Parameters:

  • metric_extractor (MetricExtractor instance): The MetricExtractor instance. This should be properly initialized and should have a model attribute.

  • dataset (Dataset instance): The dataset containing the images. This should have a file_paths attribute which is an iterable of image file paths.

  • file_path (str): The path where the plot should be saved. This should be a valid directory where the user has write permissions.

  • num_images (int, optional): The number of images to plot. Default is 5. This should be a positive integer.

Raises:

  • ValueError: If any of the inputs are not as expected.

Example usage:

from data import OptionalSplitDataset
from core.metrics import MetricExtractor

# Initialize a MetricExtractor instance
metric_extractor = MetricExtractor(model, device)

# Initialize a dataset
dataset = ImageFolder('path_to_dataset')

# Specify the file path
file_path = 'path_to_save_plot'

# Call the function
plot_activation_histograms(metric_extractor, dataset, file_path, num_images=10)
core.plotting.plot_inNout(data, model, epoch, model_name)

Plot the original images and their reconstructions side by side.

This function takes a data loader, a model, an epoch or output path, and a model name. It then plots the original images and their reconstructions side by side, and saves the plot to a specified file.

Parameters:

  • data (DataLoader instance): The DataLoader instance providing batches of images.

  • model (nn.Module instance): The PyTorch model to use for reconstruction. This should be a model that can take a batch of images and return their reconstructions.

  • epoch (int or str): The current epoch (for saving purposes), or if during metrics extraction, the desired output path.

  • model_name (str): The name of the model. This is used for saving purposes.

Raises:

  • ValueError: If any of the inputs are not as expected.

Example usage:

from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from core.model.model import Model
from core.data import OptionalSplitDataset

# Initialize a DataLoader instance
dataset = OptionalSplitDataset(root_dir='path_to_dataset', split='none', device='cuda')
data = DataLoader(dataset, batch_size=64)

# Initialize a model
model = Model(model_name='my_model', device='cuda')

# Specify the epoch and model name
epoch = 10
model_name = 'my_model'

# Call the function
plot_inNout(data, model, epoch, model_name)

core.training

This module contains the main training function.

class core.training.SAM(params, base_optimizer, rho=0.05, adaptive=False, **kwargs)

This class implements the Sharpness-aware-Minimization optimizer, taken from git.

__init__(params, base_optimizer, rho=0.05, adaptive=False, **kwargs)
__module__ = 'core.training'
first_step(zero_grad=False)
load_state_dict(state_dict)

Loads the optimizer state.

Args:
state_dict (dict): optimizer state. Should be an object returned

from a call to state_dict().

second_step(zero_grad=False)
step(closure=None)

Performs a single optimization step (parameter update).

Args:
closure (Callable): A closure that reevaluates the model and

returns the loss. Optional for most optimizers.

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

core.training.train_the_model(model, dataset, loss, opt, nb_epochs)

Train a model for a specified number of epochs.

This function takes a model, a dataset, a loss function, an optimizer, and a number of epochs, and trains the model on the dataset for the specified number of epochs. It returns the reconstruction loss and KL divergence for each batch.

Parameters:

  • model (nn.Module instance): The PyTorch model to train.

  • dataset (DataLoader instance): The DataLoader instance providing batches of images.

  • loss (nn.Module instance): The PyTorch loss function to use.

  • opt (torch.optim.Optimizer instance): The PyTorch optimizer to use.

  • nb_epochs (int): The number of epochs to train for.

  • save (int, optional): The frequency at which to save the model. Default is 10.

Returns:

  • tuple: A tuple of three lists: the reconstruction loss per batch, the KL divergence per batch, and the attention loss per batch.

Raises:

  • ValueError: If the dataset is empty.

Example usage:

from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torch import nn, optim
from core.model.models import Model
from core.data import OptionalSplitDataset

# Initialize a DataLoader instance
dataset = OptionalSplitDataset(root_dir='path_to_dataset', split='none', device='cuda')
data = DataLoader(dataset, batch_size=64)

# Initialize a model
model = Model(model_name='my_model', device='cuda')

# Specify the loss function and optimizer
loss = nn.MSELoss()
opt = optim.Adam(model.parameters())

# Specify the number of epochs and save frequency
nb_epochs = 100
save = 10

# Call the function
reconstruction_loss_per_batch, kl_divergence_per_batch, attention_per_batch = train_the_model(model, data, loss, opt, nb_epochs, save)