Skip to content

hooks

Modules

fastvideo.hooks.hooks

Classes

fastvideo.hooks.hooks.ForwardHook

Base class for forward hooks. Hooks are used in the way: modified_args, modified_kwargs = hook.pre_forward(module, args, kwargs) output = module.forward(modified_args, **modified_kwargs) modified_output = hook.post_forward(module, output)

Functions
fastvideo.hooks.hooks.ForwardHook.on_attach 
on_attach(module: Module)

Called once when the hook is attached to the module.

Source code in fastvideo/hooks/hooks.py 
21
22
23
def on_attach(self, module: nn.Module):  # noqa: B027
    """Called once when the hook is attached to the module."""
    pass
fastvideo.hooks.hooks.ForwardHook.on_detach 
on_detach(module: Module)

Called once when the hook is detached from the module. Note: this function is not guaranteed to be called if the module is deleted before the hook is detached.

Source code in fastvideo/hooks/hooks.py 
25
26
27
28
29
30
31
def on_detach(self, module: nn.Module):  # noqa: B027
    """
    Called once when the hook is detached from the module.
    Note: this function is not guaranteed to be called if the module is
          deleted before the hook is detached.
    """
    pass
fastvideo.hooks.hooks.ForwardHook.post_forward 
post_forward(module: Module, output: Any) -> Any

Called after the module's forward method is executed.

Source code in fastvideo/hooks/hooks.py 
38
39
40
def post_forward(self, module: nn.Module, output: Any) -> Any:
    """Called after the module's forward method is executed."""
    return output
fastvideo.hooks.hooks.ForwardHook.pre_forward 
pre_forward(module: Module, *args, **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]

Called before the module's forward method is executed.

Source code in fastvideo/hooks/hooks.py 
33
34
35
36
def pre_forward(self, module: nn.Module, *args,
                **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]:
    """Called before the module's forward method is executed."""
    return args, kwargs

fastvideo.hooks.layerwise_offload

Classes

fastvideo.hooks.layerwise_offload.LayerwiseOffloadHook 
LayerwiseOffloadHook(state: LayerwiseOffloadState)

Bases: ForwardHook

A hook that enables layerwise CPU offloading during forward pass.

Source code in fastvideo/hooks/layerwise_offload.py 
85
86
def __init__(self, state: LayerwiseOffloadState) -> None:
    self.state = state

Functions