yamle.methods.laplace module

class yamle.methods.laplace.LaplaceMethod(mode, *args, **kwargs)

Bases: MCSamplingMethod

This class implements the Laplace method.

It performs Laplace approximation in the neural network to approximate the posterior distribution.

It is important the regularisation is the L2Regularizer and the regularizer_weight is non-zero. These act as the prior distribution.

At the moment, it only supports the last_layer mode and the classification task. This is because the backpack library supports only the CrossEntropyLoss but not the GaussianNLLLoss.

Parameters:

mode (str) – The mode of the method. It can be last_layer for now.

tasks = ['classification']

property hessian_computed: bool

This property returns whether the Hessian is computed.

load_state_dict(state_dict)

This function loads the state dictionary of the method.

Return type:

None

on_after_model_load()

This function is called after the model is loaded.

Return type:

None

state_dict()

This function returns the state dictionary of the method.

Return type:

Dict[str, Any]

on_fit_end()

This method is used to compute the factorised Hessian at the end of the training.

It uses the K-FAC approximation to compute the factorised Hessian.

Return type:

None

static add_specific_args(parent_parser)

This function adds the arguments for the Laplace method.

Return type:

ArgumentParser