yamle.methods.sgld module

class yamle.methods.sgld.SGLDMethod(sampling_epochs=[0, 1, 2, 3, 4, 5, 10, 20, 50, 100], *args, **kwargs)

Bases: EnsembleMethod

This is a Method class for the Stochastic Gradient Langevin Dynamics Optimizer.

It uses the Ensemble model to wrap around the original model and then uses the base method to train the network via the SGLD optimizer.

At predefined epoch intervals, the weights of the main model are copied into the next member in the ensemble. This represents the posterior distribution of the weights. The last sample is always at the last epoch.

Parameters:

sampling_epochs (List[int], optional) – Epochs at which to sample (default: [0, 1, 2, 3, 4, 5, 10, 20, 50, 100]).

get_parameters(recurse=True)

A helper function to get the parameters of a single ensemble member.

In this case, get always the first one.

Return type:

List[Parameter]

on_train_epoch_end()

This method is called at the end of each training epoch.

In this case, if the current epoch can be found in the sampling epochs, the current member is incremented.

Return type:

None

static add_specific_args(parent_parser)

This method adds arguments specific to this method to the parser.

Return type:

ArgumentParser

test_name: Optional[str]

prepare_data_per_node: bool

allow_zero_length_dataloader_with_multiple_devices: bool

training: bool