torch_blue.vi.VIModule

class torch_blue.vi.VIModule(variable_shapes: Mapping[str, Tuple[int, ...] | None] | None = None, variational_distribution: Distribution | List[Distribution] = MeanFieldNormal(), prior: Distribution | List[Distribution] = MeanFieldNormal(), rescale_prior: bool = False, kaiming_initialization: bool = True, prior_initialization: bool = False, return_log_probs: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None, convert_overwrite: bool = False)

Bases: torch.nn.Module

Base class for Modules using Variational Inference.

This class takes the place of torch.nn.Module for BNNs. It is used for any Bayesian module. While a torch.nn.Module may be contained within a VIModule and vice versa there should always be a singular VIModule containing all others to avoid superfluous sampling dimensions.

VIModule contains some additional functionality. Firstly, it keeps track of whether the model should return the log probability of the sampled weight (i.e., the log probability to obtain these specific values when sampling from the prior or variational distribution).These are needed for loss calculation by KullbackLeiblerLoss. This is handled automatically as part of accessing the weight matrices, if return_log_probs is True. For evaluation, it can be set to False.

The outermost module will aggregate the log probabilities and pack the output and log probs into a VIReturn object, which behaves like a pytorch Tensor. However, it has the added attribute log_probs where the log probs are stored. This is used by losses, therefore it is easiest to wrap all operations into a VIModule and feed the output directly into a loss function. Classical losses will treat it like a Tensor and torch_blue losses can use the log prob information. If the model has multiple output Tensors, each will contain the full log prob information.

Important

When defining custom modules with weights make sure to retrieve them using sample_variables() as this will maintain the automatic log prob tracking.

Secondly, a model of nested VIModule automatically identifies the outermost module and sets the _has_sampling_responsibility flag. This makes the outermost module always accept the keyword argument samples, which defaults to 10. Since BNNs require multiple samples for each forward pass, the input batch is duplicated accordingly and the forward pass is performed vectorized on all samples. While this is significantly faster than serial evaluation, it naturally requires more memory. Additionally, a certain few operations do not function correctly with the vectorization and should not be used in VIModule. Most importantly, this affects the operators +=, -=, *=, and /=. However, their longform versions work fine, e.g. a = a + b instead of a += b.

If the constructed module does not have its own weights, super().__init__() is called without arguments. In this setting the methods get_log_probs(), get_variational_parameters(), reset_variational_parameters(), and sample_variable() cannot be used and raise NoVariablesError.

Any weight matrix in a BNN may require multiple parameters (e.g. mean and std). These are stored in separate attributes and therefore cannot be created without knowledge of the variational distribution. Therefore, only their shapes are provided to super().__init__() via the variable_shapes argument. This should be a dictionary, where the keys are the random variable names as strings (e.g. “weight” and “bias”) and the values are tuples of integers specifying the shape.

After initialization these variables can be accessed as an attribute under the specified name. Each access will yield a new sample. The shape of a random variable can be set to None. In that case accessing it will always return None.

Note

The insertion order of the dictionary becomes the order self.random_variables.

The names of the created attributes can be discovered using the variational_parameter_name() method.

Additionally, a module with random variables accepts arguments from VIkwargs as keyword arguments. If a list of priors or variational distributions is provided they are again assumed to follow the insertion order as described above.

Parameters:

• variable_shapes (Optional[Mapping[str, Optional[Tuple[int, ...]]]], default = None) – Shape specifications for all random variables. Keys are turned into self.random_variables in insertion order.

• VIkwargs – Several standard keyword arguments. See VIkwargs for details.

Raises:

NoVariablesError – If the module does not have Bayesian parameters of its own and a method requiring them is called.

property random_variables: Tuple[str, Ellipsis] | None

Names of the modules random variables.

reset_variational_parameters() → None

Reset or initialize the parameters of the Module.

Raises:

NoVariablesError – If the module does not have parameters of its own.

static variational_parameter_name(variable: str, variational_parameter: str) → str

Get the attribute name of the variational parameter for the specified variable.

Parameters:

• variable (str) – Random variable name as specified in self.random_variables.

• variational_parameter (str) – Variational parameter name as specified by the variational distribution.

Returns:

The attribute name of the specified parameter.

Return type:

str

get_variational_parameters(variable: str) → List[torch.Tensor]

Get all variational parameters for the specified variable.

Parameters:

variable (str) – Random variable name as specified in self.random_variables.

Returns:

All variational parameters for the specified variable in the order specified by the associated variational distribution.

Return type:

List[Tensor]

Raises:

NoVariablesError – If the module does not have parameters of its own or the requested variable is None.

get_log_probs(sample: torch.Tensor, variable: str) → torch.Tensor

Get prior and variational log prob of the sampled parameters.

Accepts the sampled parameters as returned by the self.sample_variable method and calculates the total prior and variational log probability.

Parameters:

• sample (Tensor) – Sampled parameter as returned by the self.sample_variable method.

• variable (str) – The name of the relevant variable.

Returns:

A Tensor containing two values: the log probability of the sampled values, if they were drawn from the prior or variational distribution (in that order).

Return type:

Tensor

Raises:

NoVariablesError – If the module does not have parameters of its own.

sample_variable(variable: str) → torch.Tensor | None

Draw one sample from the variational distribution of one random variable.

This also performs log prob tracking, if return_log_probs is True.

Parameters:

variable (str) – The variable to sample.

Returns:

The sampled variable.

Return type:

Tensor

Raises:

NoVariablesError – If the module does not have parameters of its own.

sampled_forward(*input_: torch.Tensor | None, samples: int = 10, **kwargs: Any) → torch_blue.vi.utils.vi_return.VIReturn | Tuple[torch_blue.vi.utils.vi_return.VIReturn, ...]

Forward pass of the module evaluating multiple weight samples.

This will automatically be called by the outermost module. Instead of the forward() method. It grabs the samples argument, if provided, and copies the input batch the specified number of times. The forward() is performed vectorized over that additional sample dimension.

If you need more information on how this is achieved, check the documentation of this class’ __post_init__ method in the source code. Note this goes in deep and if checking the source code seems too much hassle you should probably not touch it and write an issue instead.

Parameters:

• input (Tensor) – Any number of input Tensors

• samples (int, default: 10) – Number of weight samples to evaluate

• kwargs (Any) – Any additional keyword arguments

Returns:

One or multiple Tensors with log prob annotation

Return type:

Union[VIReturn, Tuple[VIReturn, …]]

reset_log_probs() → None

Reset all tracked log probabilities.

gather_log_probs() → torch.Tensor

Gather and aggregate log probs from all submodules, then reset them.

Returns:

The aggregated log probabilities of all submodules.

Return type:

Tensor

property return_log_probs: bool

Set whether the module returns log probabilities.

Log probabilities are required for most standard losses.

Returns:

Whether the module returns log probabilities.

Return type:

bool