yamle.methods.sngp module

yamle.methods.sngp.enable_spectral_normalization(model, coeff)

Replace all the layers in the model with spectral normalized layers.

Parameters:

model (torch.nn.Module) – The model to enable spectral normalization for.

Return type:

None

class yamle.methods.sngp.SNGPMethod(m=0.99, random_features=512, mean_field_factor=1.0, coeff=1.0, **kwargs)

Bases: BaseMethod

This class is the extension of the base method for which the prediciton is performed through the method of: Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness

The core of the method is to 1. enable spectral normalization for all ._residual layers in the model and replace the ._output layer with a ._output layer with a Gaussian process

Parameters:

• m (float) – The gamma for exponential moving average for updating the precision matrix.

• random_features (int) – The number of random features to use in the RFF layer.

• mean_field_factor (float) – The factor to use for the mean field approximation.

• coeff (float) – The coefficient for the spectral normalization.

tasks = ['classification', 'segmentation']

on_train_epoch_start()

In the final epoch we need to update the precision matrix. The update is triggered by the _final_epoch flag set to True.

Return type:

None

on_train_epoch_end()

This method is used to:

Reset the model at the end of each training epoch. Step the learning rate schedulers if automatic optimization is not selected. Plot the training results if plotting is selected. Apply the regularizer at the end of each training epoch if a regularizer is selected.

Return type:

None

static add_specific_args(parent_parser)

This method is used to add the specific arguments for the DUN method.

Return type:

ArgumentParser

test_name: Optional[str]

prepare_data_per_node: bool

allow_zero_length_dataloader_with_multiple_devices: bool

training: bool