hooks #

Module hooks for forward and backward inspection (issue #808).

with_hooks wraps any types.Layer without changing the layer or the Layer type. The wrapper is itself a types.Layer, so it can be placed in a Sequential or nested inside another wrapper. The registry is not part of the layer's variables or serialized parameters.

Contract:- Forward. Pre-hooks run first, in registration order. Each receives the input produced by the previous hook and may replace it. The wrapped layer then runs. Post-hooks run next, in registration order, each receiving the input the layer saw and the output of the previous hook. A hook error stops the forward pass and is returned unchanged.- Backward. Backward hooks are attached to the module output only when that output requires a gradient. They run in registration order, each receiving the gradient produced by the previous hook, and the last result flows into the layer output. A hook error stops backpropagation and is returned from the backward call.- No gradients. Backward hooks attach only when Variable.is_grad_needed() holds, that is requires_grad is set and the context is not no_grad. Pre- and post-hooks run either way.- Removal. Each registration returns a HookHandle. remove() takes the hook out of the registry, so later passes do not run it. Calling remove() again has no effect.- Order between kinds: forward_pre, the layer, forward, then backward in reverse graph order.

fn with_hooks[T](layer types.Layer[T]) (types.Layer[T], Hooks[T])

with_hooks wraps layer. It returns the wrapper, used in place of layer, and the Hooks object used to register hooks on it.

fn HookKind.from[W](input W) !HookKind

from returns the HookKind value of an integer or of a field name, or an error if there is none.

type BackwardHook[T] = fn (grad &vtl.Tensor[T]) !&vtl.Tensor[T]

BackwardHook receives the gradient with respect to the module output and may replace it.

fn (g &BackwardHookGate[T]) backward(payload &autograd.Payload[T]) ![]&vtl.Tensor[T]

backward applies each registered backward hook to the gradient, in registration order.

fn (g &BackwardHookGate[T]) cache(mut result autograd.Variable[T], args ...autograd.CacheParam) !

cache records the gate in the autograd graph, with the module output as its only parent.

type ForwardHook[T] = fn (input &autograd.Variable[T], output &autograd.Variable[T]) !&autograd.Variable[T]

ForwardHook receives the input the layer saw and the output, and may replace the output.

type ForwardPreHook[T] = fn (input &autograd.Variable[T]) !&autograd.Variable[T]

ForwardPreHook may replace the module input.

fn (h HookHandle[T]) remove()

remove unregisters the hook. Calling it again has no effect.

fn (h Hooks[T]) add_forward_pre_hook(hook ForwardPreHook[T]) HookHandle[T]

add_forward_pre_hook registers a pre-hook and returns its handle.

fn (h Hooks[T]) add_forward_hook(hook ForwardHook[T]) HookHandle[T]

add_forward_hook registers a post-hook and returns its handle.

fn (h Hooks[T]) add_backward_hook(hook BackwardHook[T]) HookHandle[T]

add_backward_hook registers a backward hook and returns its handle.

enum HookKind {
	forward_pre
	forward
	backward
}

HookKind names the three kinds of registration.

struct HookHandle[T] {
	registry &HookRegistry[T] = unsafe { nil }
	kind     HookKind
	id       int
}

HookHandle removes one registration.

struct HookRegistry[T] {
mut:
	next_id int
	pre     []PreEntry[T]
	post    []PostEntry[T]
	back    []BackEntry[T]
}

HookRegistry holds the registered hooks of one wrapped module. The wrapper and every HookHandle share it, so removal is visible to later passes.

struct Hooks[T] {
pub:
	registry &HookRegistry[T]
}

Hooks gives access to the registration methods of a hooked wrapper.