Table of Contents
Shortcuts

torchfunc.hooks.registrators

This module allows you for easier hook registration (e.g. based on type or index within network).

Example:

# Example forward pre hook
def example_forward_pre(module, inputs):
    return inputs + 1

# MNIST classifier
model = torch.nn.Sequential(
    torch.nn.Linear(784, 100),
    torch.nn.ReLU(),
    torch.nn.Linear(100, 50),
    torch.nn.ReLU(),
    torch.nn.Linear(50, 10),
)
registrator = torchfunc.hooks.registrators.ForwardPre()
# Register forwardPreHook for all torch.nn.Linear submodules
registrator.modules(model, example_forward_pre, types=(torch.nn.Linear))

You could specify indices instead of types (for example all inputs to torch.nn.Linear will be registered), and iterate over children instead of modules.

class torchfunc.hooks.registrators.Backward(hook: Callable)[source]

Register backward hook based on module’s type or indices.

handles

Handles for registered hooks, each corresponds to specific submodule. Can be used to unregister certain hooks (though discouraged).

Type

List[torch.utils.hooks.RemovableHandle]

children(network, types: Tuple[Any] = None, indices: List[int] = None)

Register subrecorders using types and/or indices via children hook.

This function will use children method of torch.nn.Module to iterate over available submodules. If you wish to iterate recursively, use modules.

Important:

If types and indices are left with their default values, all modules will have subrecorders registered.

Parameters

  • module (torch.nn.Module) – Module (usually neural network) for which inputs will be collected.

  • types (Tuple[typing.Any], optional) – Module types for which data will be recorded. E.g. (torch.nn.Conv2d, torch.nn.Linear) will register subrecorders on every module being instance of either Conv2d or Linear. Default: None

  • indices (Iterable[int], optional) – Indices of modules whose inputs will be registered. Default: None

Returns

Return type

self

modules(module: torch.nn.modules.module.Module, types: Tuple[Any] = None, indices: List[int] = None)

Register hook using types and/or indices via modules hook.

This function will use modules method of torch.nn.Module to iterate over available submodules. If you wish to iterate non-recursively, use children.

Important:

If types and indices are left with their default values, all modules will have subrecorders registered.

Parameters

  • module (torch.nn.Module) – Module (usually neural network) for which inputs will be collected.

  • types (Tuple[typing.Any], optional) – Module types for which data will be recorded. E.g. (torch.nn.Conv2d, torch.nn.Linear) will register subrecorders on every module being instance of either Conv2d or Linear. Default: None

  • indices (Iterable[int], optional) – Indices of modules whose inputs will be registered. Default: None

Returns

Return type

self

remove(index) → None

Remove hook specified by index.

Parameters

index (int) – Index of subhook (usually registered for layer)

class torchfunc.hooks.registrators.Forward(hook: Callable)[source]

Register forward hook based on module’s type or indices.

handles

Handles for registered hooks, each corresponds to specific submodule. Can be used to unregister certain hooks (though discouraged).

Type

List[torch.utils.hooks.RemovableHandle]

children(network, types: Tuple[Any] = None, indices: List[int] = None)

Register subrecorders using types and/or indices via children hook.

This function will use children method of torch.nn.Module to iterate over available submodules. If you wish to iterate recursively, use modules.

Important:

If types and indices are left with their default values, all modules will have subrecorders registered.

Parameters

  • module (torch.nn.Module) – Module (usually neural network) for which inputs will be collected.

  • types (Tuple[typing.Any], optional) – Module types for which data will be recorded. E.g. (torch.nn.Conv2d, torch.nn.Linear) will register subrecorders on every module being instance of either Conv2d or Linear. Default: None

  • indices (Iterable[int], optional) – Indices of modules whose inputs will be registered. Default: None

Returns

Return type

self

modules(module: torch.nn.modules.module.Module, types: Tuple[Any] = None, indices: List[int] = None)

Register hook using types and/or indices via modules hook.

This function will use modules method of torch.nn.Module to iterate over available submodules. If you wish to iterate non-recursively, use children.

Important:

If types and indices are left with their default values, all modules will have subrecorders registered.

Parameters

  • module (torch.nn.Module) – Module (usually neural network) for which inputs will be collected.

  • types (Tuple[typing.Any], optional) – Module types for which data will be recorded. E.g. (torch.nn.Conv2d, torch.nn.Linear) will register subrecorders on every module being instance of either Conv2d or Linear. Default: None

  • indices (Iterable[int], optional) – Indices of modules whose inputs will be registered. Default: None

Returns

Return type

self

remove(index) → None

Remove hook specified by index.

Parameters

index (int) – Index of subhook (usually registered for layer)

class torchfunc.hooks.registrators.ForwardPre(hook: Callable)[source]

Register forward pre hook based on module’s type or indices.

handles

Handles for registered hooks, each corresponds to specific submodule. Can be used to unregister certain hooks (though discouraged).

Type

List[torch.utils.hooks.RemovableHandle]

children(network, types: Tuple[Any] = None, indices: List[int] = None)

Register subrecorders using types and/or indices via children hook.

This function will use children method of torch.nn.Module to iterate over available submodules. If you wish to iterate recursively, use modules.

Important:

If types and indices are left with their default values, all modules will have subrecorders registered.

Parameters

  • module (torch.nn.Module) – Module (usually neural network) for which inputs will be collected.

  • types (Tuple[typing.Any], optional) – Module types for which data will be recorded. E.g. (torch.nn.Conv2d, torch.nn.Linear) will register subrecorders on every module being instance of either Conv2d or Linear. Default: None

  • indices (Iterable[int], optional) – Indices of modules whose inputs will be registered. Default: None

Returns

Return type

self

modules(module: torch.nn.modules.module.Module, types: Tuple[Any] = None, indices: List[int] = None)

Register hook using types and/or indices via modules hook.

This function will use modules method of torch.nn.Module to iterate over available submodules. If you wish to iterate non-recursively, use children.

Important:

If types and indices are left with their default values, all modules will have subrecorders registered.

Parameters

  • module (torch.nn.Module) – Module (usually neural network) for which inputs will be collected.

  • types (Tuple[typing.Any], optional) – Module types for which data will be recorded. E.g. (torch.nn.Conv2d, torch.nn.Linear) will register subrecorders on every module being instance of either Conv2d or Linear. Default: None

  • indices (Iterable[int], optional) – Indices of modules whose inputs will be registered. Default: None

Returns

Return type

self

remove(index) → None

Remove hook specified by index.

Parameters

index (int) – Index of subhook (usually registered for layer)