TF-MELT package

TF-MELT is composed of a series of modules that form the basis for building a variety of machine learning models. The structure of the packages is to encourage modularity and reusability of the code. The main modules are:

  • Blocks Module: Each of the blocks is a self-contained TensorFlow model in itself, but can be combined with other blocks to form more complex models. The blocks are designed to be easily combined with other blocks, and to be easily extended.

  • Losses Module: The losses are custom loss functions that should be used with their respective models. Certain loss functions are designed to be used with specific models, but others can be used with any model.

  • Models Module: The models are the main machine learning models that are built using the blocks and losses. The models serve a dual purpose of being a standalone model, and also as a template for building more complex models with the TF-MELT blocks.

Following is a detailed description of the modules and subpackages in the TF-MELT.

Blocks Module

class tfmelt.blocks.BayesianAleatoricOutput(*args, **kwargs)[source]

Bases: Model

Output layer for a Bayesian neural network with aleatoric uncertainty.

Parameters:

  • num_outputs (int) – Number of output nodes.

  • num_points (int) – Number of Monte Carlo samples.

  • regularizer (Regularizer, optional) – Kernel regularizer. Defaults to None.

  • scale_epsilon (float, optional) – Epsilon value for scale parameter. Defaults to 1e-3.

  • aleatoric_scale_factor (float, optional) – Scaling factor for aleatoric uncertainty. Defaults to 5e-2.

  • **kwargs – Extra arguments passed to the base class.

call(x, training=False)[source]

Forward pass through the output layer.

classmethod from_config(config)[source]

Create model from config dictionary

get_config()[source]

Get the config dictionary

class tfmelt.blocks.BayesianBlock(*args, **kwargs)[source]

Bases: MELTBlock

A BayesianBlock consists of multiple Bayesian dense layers with optional activation, dropout, and batch normalization. The layers are implemented using the Flipout variational layer.

Parameters:

  • num_points (int, optional) – Number of Monte Carlo samples. Defaults to 1.

  • use_batch_renorm (bool, optional) – Use batch renormalization. Defaults to True.

  • **kwargs – Extra arguments passed to the base class.

call(inputs: Tensor, training: bool = False)[source]

Forward pass through the Bayesian block.

class tfmelt.blocks.DefaultOutput(*args, **kwargs)[source]

Bases: Model

Default output layer with a single dense layer.

Parameters:

  • num_outputs (int) – Number of output nodes.

  • output_activation (str, optional) – Activation function for the output layer. Defaults to None.

  • initializer (str, optional) – Kernel initializer. Defaults to “glorot_uniform”.

  • regularizer (Regularizer, optional) – Kernel regularizer. Defaults to None.

  • bayesian (bool, optional) – Use Bayesian layer if True. Defaults to False.

  • num_points (int, optional) – Number of samples. Defaults to 1.

  • **kwargs – Extra arguments passed to the base class.

call(x, training=False)[source]

Forward pass through the output layer.

classmethod from_config(config)[source]

Create model from config dictionary

get_config()[source]

Get the config dictionary

class tfmelt.blocks.DenseBlock(*args, **kwargs)[source]

Bases: MELTBlock

A DenseBlock consists of multiple dense layers with optional activation, dropout, and batch normalization.

Parameters:

**kwargs – Extra arguments passed to the base class.

Raises:

AssertionError – If dropout is not within the range of [0, 1].

call(inputs: Tensor, training: bool = False)[source]

Forward pass through the dense block.

class tfmelt.blocks.MELTBlock(*args, **kwargs)[source]

Bases: Model

Base class for a MELT block. Defines a block of dense layers with optional activation, dropout, and batch normalization. Forms the building block for various neural network blocks.

Parameters:

  • node_list (List[int]) – Number of nodes in each dense layer. The length of the list determines the number of layers.

  • activation (str, optional) – Activation function. If None, no activation is applied (linear). Defaults to “relu”.

  • dropout (float, optional) – Dropout rate (0-1). Defaults to None.

  • batch_norm (bool, optional) – Apply batch normalization if True. Defaults to False.

  • use_batch_renorm (bool, optional) – Use batch renormalization. Defaults to False.

  • regularizer (Regularizer, optional) – Kernel weights regularizer. Defaults to None.

  • initializer (str, optional) – String defining the kernel initializer. Defaults to “glorot_uniform”.

  • seed (int, optional) – Random seed. Defaults to None.

  • **kwargs – Extra arguments passed to the base class.

classmethod from_config(config)[source]

Create model from config dictionary

get_config()[source]

Get the config dictionary

class tfmelt.blocks.MixtureDensityOutput(*args, **kwargs)[source]

Bases: Model

Output layer for a mixture density network.

Parameters:

  • num_mixtures (int) – Number of mixture components.

  • num_outputs (int) – Number of output nodes.

  • output_activation (str, optional) – Activation function for the output layer. Defaults to None.

  • initializer (str, optional) – Kernel initializer. Defaults to “glorot_uniform”.

  • regularizer (Regularizer, optional) – Kernel regularizer. Defaults to None.

  • **kwargs – Extra arguments passed to the base class.

call(x, training=False)[source]

Forward pass through the output layer.

classmethod from_config(config)[source]

Create model from config dictionary

get_config()[source]

Get the config dictionary

class tfmelt.blocks.ResidualBlock(*args, **kwargs)[source]

Bases: MELTBlock

A ResidualBlock consists of multiple dense layers with optional activation, dropout, and batch normalization. Residual connections are added after each block of layers.

Parameters:

  • layers_per_block (int, optional) – Number of layers per residual block. Defaults to 2.

  • pre_activation (bool, optional) – Apply activation before adding residual connection. Defaults to False.

  • post_add_activation (bool, optional) – Apply activation after adding residual connection. Defaults to False.

  • **kwargs – Extra arguments passed to the base class.

call(inputs: Tensor, training: bool = False)[source]

Forward pass through the residual block.

tfmelt.blocks.get_kernel_divergence_fn(num_points)[source]

Get the kernel divergence function.

Parameters:

num_points (int) – Number of points in kernel divergence.

Returns:

Kernel divergence function.

Return type:

Callable

Losses Module

class tfmelt.losses.MixtureDensityLoss(num_mixtures, num_outputs, **kwargs)[source]

Bases: Loss

Loss function for the Mixture Density Network (MDN) model. Computes the negative log likelihood using the weighted average of the Gaussian mixture model components.

Parameters:

  • num_mixtures (int) – Number of mixture components.

  • num_outputs (int) – Number of output dimensions.

call(y_true, y_pred)[source]

Invokes the Loss instance.

Parameters:

  • y_true – Ground truth values. shape = [batch_size, d0, .. dN], except sparse loss functions such as sparse categorical crossentropy where shape = [batch_size, d0, .. dN-1]

  • y_pred – The predicted values. shape = [batch_size, d0, .. dN]

Returns:

Loss values with the shape [batch_size, d0, .. dN-1].

get_config()[source]

Get the config dictionary.

tfmelt.losses.safe_exp(x)[source]

Prevents overflow by clipping input range to reasonable values.

The function clips the input range and then takes the exponential:

\[\text{safe_exp}(x) = \exp(\min(\max(x, -20), 20))\]
Parameters:

x (tensor) – Input tensor.

Models Module

class tfmelt.models.ArtificialNeuralNetwork(*args, **kwargs)[source]

Bases: MELTModel

Artificial Neural Network model.

Parameters:

**kwargs – Additional keyword arguments.

build(input_shape)[source]

Build the ANN.

call(inputs, training=False)[source]

Call the ANN.

get_config()[source]

Get the config dictionary.

initialize_layers()[source]

Initialize the layers of the ANN.

class tfmelt.models.BayesianNeuralNetwork(*args, **kwargs)[source]

Bases: MELTModel

Bayesian Neural Network model.

Parameters:

  • num_points (int, optional) – Number of Monte Carlo samples. Defaults to 1.

  • do_aleatoric (bool, optional) – Flag to perform aleatoric output. Defaults to False.

  • do_bayesian_output (bool, optional) – Flag to perform Bayesian output. Defaults to True.

  • aleatoric_scale_factor (float, optional) – Scale factor for aleatoric uncertainty. Defaults to 5e-2.

  • scale_epsilon (float, optional) – Epsilon value for the scale of the aleatoric uncertainty. Defaults to 1e-3.

  • use_batch_renorm (bool, optional) – Whether to use batch renormalization. Defaults to True.

  • bayesian_mask (list, optional) – List of booleans to determine which layers are Bayesian and which are Dense. Defaults to None.

  • **kwargs – Additional keyword arguments.

build(input_shape)[source]

Build the BNN.

call(inputs, training=False)[source]

Call the BNN.

compile(optimizer='adam', loss='mse', metrics=None, **kwargs)[source]

Compile the model with the appropriate loss function.

create_output_layer()[source]

Create output layer for the Bayesian Neural Network.

get_config()[source]

Get the config dictionary.

initialize_layers()[source]

Initialize the layers of the BNN.

negative_log_likelihood(y_true, y_pred)[source]

Calculate the negative log likelihood.

class tfmelt.models.MELTModel(*args, **kwargs)[source]

Bases: Model

TF-MELT base model.

Parameters:

  • num_outputs (int) – Number of output units.

  • width (int, optional) – Width of the hidden layers. Defaults to 32.

  • depth (int, optional) – Number of hidden layers. Defaults to 2.

  • act_fun (str, optional) – Activation function for the hidden layers. Defaults to “relu”.

  • dropout (float, optional) – Dropout rate for the hidden layers. Defaults to 0.0.

  • input_dropout (float, optional) – Dropout rate for the input layer. Defaults to 0.0.

  • batch_norm (bool, optional) – Whether to use batch normalization. Defaults to False.

  • use_batch_renorm (bool, optional) – Whether to use batch renormalization. Defaults to False.

  • output_activation (str, optional) – Activation function for the output layer. Defaults to None.

  • initializer (str, optional) – Initializer for the weights. Defaults to “glorot_uniform”.

  • seed (int, optional) – Random seed for reproducibility. Defaults to None.

  • l1_reg (float, optional) – L1 regularization for the weights. Defaults to 0.0.

  • l2_reg (float, optional) – L2 regularization for the weights. Defaults to 0.0.

  • num_mixtures (int, optional) – Number of mixtures for density networks. Defaults to 0.

  • node_list (list, optional) – Numbers of nodes to alternately define layers. Defaults to None.

  • **kwargs – Additional keyword arguments.

compile(optimizer='adam', loss='mse', metrics=None, **kwargs)[source]

Compile the model with the appropriate loss function.

compute_jacobian(x)[source]

Compute the Jacobian of the model outputs with respect to inputs.

create_dropout_layers()[source]

Create the dropout layers.

create_output_layer()[source]

Create the output layer based on the number of mixtures.

create_regularizer()[source]

Create the regularizer.

classmethod from_config(config)[source]

Create a model from a config dictionary.

get_config()[source]

Get the config dictionary.

initialize_layers()[source]

Initialize the layers of the model.

summary()[source]

Print a summary of the model that includes sub-layers.

class tfmelt.models.ResidualNeuralNetwork(*args, **kwargs)[source]

Bases: MELTModel

Residual Neural Network model.

Parameters:

  • layers_per_block (int, optional) – Number of layers in each block. Defaults to 2.

  • pre_activation (bool, optional) – Whether to use pre-activation in residual blocks. Defaults to True.

  • post_add_activation (bool, optional) – Whether to apply activation after adding the residual connection. Defaults to False.

  • **kwargs – Additional keyword arguments.

build(input_shape)[source]

Build the ResNet.

call(inputs, training=False)[source]

Call the ResNet.

get_config()[source]

Get the config dictionary.

initialize_layers()[source]

Initialize the layers of the ResNet.

Subpackages

In addition to the main modules, there are subpackages that contain various utility functions that are used in the main modules. The subpackages are:

  • TF-MELT UtilitiesContains utility functions that are used in the main modules. These

    functions contain routines for data processing, model evaluation, visualization, and other general-purpose functions.