def register_intern_hook(self, hook: ActionHook) -> RemovableHandle:
"""Registers an intern hook on the exporter.
The hook will be called each time a module matches against an :meth:`intern` pattern.
It should have the following signature::
hook(exporter: PackageExporter, module_name: str) -> None
Hooks will be called in order of registration.
Returns:
:class:`torch.utils.hooks.RemovableHandle`:
A handle that can be used to remove the added hook by calling
``handle.remove()``.
"""
handle = RemovableHandle(self._intern_hooks)
self._intern_hooks[handle.id] = hook
return handle
def register_batch_start_hook(
self, hook: Callable[[int, int, int, Any],
None]) -> RemovableHandle:
"""
Called at the start of a batch with the following info:
(counter, step_count, batch_size, data)
where counter is passed in to the run (ex: epoch),
step_count is the number of items run so far,
batch_size is the number of elements fed in the batch,
data is the data output from the loader
:param hook: the hook to add that is called into when reached in the
batch process
:return: a removable handle to remove the hook when desired
"""
handle = RemovableHandle(self._batch_start_hooks)
self._batch_start_hooks[handle.id] = hook
return handle
def register_propagate_forward_pre_hook(self,
hook: Callable) -> RemovableHandle:
r"""Registers a forward pre-hook on the module.
The hook will be called every time before :meth:`propagate` is invoked.
It should have the following signature:
.. code-block:: python
hook(module, inputs) -> None or modified input
The hook can modify the input.
Input keyword arguments are passed to the hook as a dictionary in
:obj:`inputs[-1]`.
Returns a :class:`torch.utils.hooks.RemovableHandle` that can be used
to remove the added hook by calling :obj:`handle.remove()`.
"""
handle = RemovableHandle(self._propagate_forward_pre_hooks)
self._propagate_forward_pre_hooks[handle.id] = hook
return handle
def register_batch_end_hook(
self, hook: Callable[[int, int, int, Any, Any, Dict[str, Tensor]],
None]):
"""
Called after all calculations are done for the batch with the following info:
(counter, step_count, batch_size, data, pred, losses)
where counter is passed in to the run (ex: epoch),
step_count is the number of items run so far,
batch_size is the number of elements fed in the batch,
data is the data output from the loader,
pred is the result from the model after the forward,
losses are the resulting loss dictionary
:param hook: the hook to add that is called into when reached in the
batch process
:return: a removable handle to remove the hook when desired
"""
handle = RemovableHandle(self._batch_end_hooks)
self._batch_end_hooks[handle.id] = hook
return handle
