pruning_mixin

Classes

LayerDescriptor
PruningMixIn

class LayerDescriptor

Bases: object

block_idx_from_module_name(module_name)

Parameters:

module_name (str)

Return type:

int | None

get_modules_names_to_hook(model)

Return type:

List[Tuple[int, str]]

module_name_regex()

Return type:

str

class PruningMixIn

Bases: ABC

__init__(layer_descriptor)

Parameters:

layer_descriptor (LayerDescriptor)

get_module_names_to_hook(model)

Return type:

List[Tuple[int, str]]

abstract supported_hooks()

Return type:

List[Type[ForwardHook]]