utils

Provides a function to register activation hooks for a model. Activation hooks are used to compute activation scores for pruning.

Functions

register_activation_hooks

Register activation hooks using the pruning mixin approach.
register_activation_hooks(model, activation_hooks_kwargs, pruning_mixin, hook_class)

Register activation hooks using the pruning mixin approach.

Parameters:

  • model – The model to register hooks on.

  • activation_hooks_kwargs (dict) – Keyword arguments passed to hook constructors.

  • pruning_mixin – The pruning mixin that defines which modules to hook.

  • hook_class (Type[ForwardHook]) – The hook class to instantiate for each module.

Returns:

Dictionary mapping module names to hook instances.

Return type:

dict[str, ForwardHook]