yamle.methods.pe module#

yamle.methods.pe.replace_layers_with_grouped_convs(model, M, alpha, gamma)[source]#

This method replaces all nn.Linear and nn.Conv2d layers with grouped versions.

Each layer is replaced with a layer where the input and output dimensions are multiplied by the alpha.

Parameters:

  • model (nn.Module) – The model to replace the layers of.

  • M (int) – The number of members in the ensemble.

  • alpha (int) – The width multiplier.

  • gamma (int) – The subgroups multiplier.

Return type:

Module

class yamle.methods.pe.PEMethod(alpha, gamma, *args, **kwargs)[source]#

Bases: MIMOMethod

This class is the extension of the base method for packed-ensemble methods.

Parameters:

  • alpha (int) – The expansion multiplier for the width of the model. It is used to multiply the number of input and output channels of each layer.

  • gamma (int) – The subgroups multiplier. It is used to multiply the number of groups in each convolutional layer together with the num_members parameter, such as (num_members * gamma).

analyse(save_path)[source]#

This method analyses the model and saves the results to a file.

Return type:

None

static add_specific_args(parent_parser)[source]#

This method adds the specific arguments for the MIMO method.

Return type:

ArgumentParser

test_name: Optional[str]#
prepare_data_per_node: bool#
allow_zero_length_dataloader_with_multiple_devices: bool#
training: bool#