fastvideo.v1.training.activation_checkpoint

fastvideo.v1.training.activation_checkpoint#

Module Contents#

Classes#

CheckpointType

Functions#

apply_activation_checkpointing

Data#

TRANSFORMER_BLOCK_NAMES

API#

class fastvideo.v1.training.activation_checkpoint.CheckpointType[source]#

Bases: str, enum.Enum

BLOCK_SKIP[source]#

‘block_skip’

FULL[source]#

‘full’

OPS[source]#

‘ops’

fastvideo.v1.training.activation_checkpoint.TRANSFORMER_BLOCK_NAMES[source]#

[‘blocks’, ‘double_blocks’, ‘single_blocks’, ‘transformer_blocks’, ‘temporal_transformer_blocks’, ‘t…

fastvideo.v1.training.activation_checkpoint.apply_activation_checkpointing(module: torch.nn.Module, checkpointing_type: str = CheckpointType.FULL, n_layer: int = 1) → torch.nn.Module[source]#
Contents