nntoolbox.hooks.io module

class nntoolbox.hooks.io.InputHook(module: torch.nn.modules.module.Module, forward: bool = True)

Bases: nntoolbox.hooks.hooks.Hook

Keep this for backward compatibility

class nntoolbox.hooks.io.OutputHook(module: torch.nn.modules.module.Module, forward: bool = True)

Bases: nntoolbox.hooks.hooks.Hook

A generic hook for storing output hook

A subclass would implement the store_output function (i.e decide which part of the output to store) (UNTESTED)

store: torch.Tensor

static store_output(hook: nntoolbox.hooks.hooks.Hook, m: torch.nn.modules.module.Module, inp: torch.Tensor, op: torch.Tensor)