norse.torch.utils.tensorboard.hook_spike_image

norse.torch.utils.tensorboard.hook_spike_image(key: str, writer: SummaryWriter) → Callable[[Module, Tuple[Tensor, Any], Any], None]

Generates a hook that can be applied to stateful torch Modules. The hook plots the spiked output, assuming that the module output is a tuple of (spikes, state). That is, this will not work on modules that only returns a single tensor as output (i. e. have no state).

Read more about hooks in PyTorch Modules in the Module documentation.

Example:

>>> from torch.utils.tensorboard import SummaryWriter
>>> import norse.torch as snn
>>> from snn.util import tensorboard
>>> hook = tensorboard.hook_spike_image("lif", SummaryWriter())
>>> snn.LIFCell().register_forward_hook(hook)

Arguments:

key (str): The name of the module as referred to in Tensorboard writer (SummaryWriter): The Tensorboard Writer to log to.

Returns:

A hook that will plot the spike activity as an image when registered to a Torch module.