neural_compressor.compression.pruner.wanda.utils

Wanda utils.

Functions

find_layers(module[, op_types, name])

Recursively find the layers of a certain type in a module.

Module Contents

neural_compressor.compression.pruner.wanda.utils.find_layers(module, op_types=['Linear', 'Conv1D'], name='')[source]

Recursively find the layers of a certain type in a module.

Parameters:

  • module (nn.Module) – PyTorch module.

  • layers (list) – List of layer types to find.

  • name (str) – Name of the module.

Returns:

Dictionary of layers of the given type(s) within the module.

Return type:

dict