modelzoo.common.pytorch.dump_context.DumpContext#

class modelzoo.common.pytorch.dump_context.DumpContext[source]#

Bases: contextlib.ContextDecorator

A debug utility context manager. When provided with a torch.nn.Module, the resulting context manager can be entered to enable dumping of all module forward and backward outputs to a npz, for comparing numerics between implementations.

Sets up global module hoooks to either dump intermediate activations on CPU/GPU or name the traced tensors for correlating with debug dumps on CS2.

The recursive name of the torch.nn.Module is memoized, and the output of FWD and BWD of each module is saved as keys in a .npz file.

Parameters

  • outdir – Where to output dumps_{i}.npz

  • model – root module to name its children

  • buffer_steps – If given, flush to a new .npz file after this many steps

Methods

disable_collection

Uninstall the hooks installed during enable_collection, disabling further dump collection.

enable_collection

Install the hooks defined during setup_hooks, enabling the collection of the dumps.

flush

Write all dump buffers out to disk.

setup_hooks

Define hooking functions on the given torch.nn.Module, but don't install them.
__call__(func)#

Call self as a function.

__init__(outdir: str, model: torch.nn.Module, buffer_steps: Optional[int] = None)[source]#

Sets up global module hoooks to either dump intermediate activations on CPU/GPU or name the traced tensors for correlating with debug dumps on CS2.

The recursive name of the torch.nn.Module is memoized, and the output of FWD and BWD of each module is saved as keys in a .npz file.

Parameters

  • outdir – Where to output dumps_{i}.npz

  • model – root module to name its children

  • buffer_steps – If given, flush to a new .npz file after this many steps

disable_collection()[source]#

Uninstall the hooks installed during enable_collection, disabling further dump collection.

enable_collection()[source]#

Install the hooks defined during setup_hooks, enabling the collection of the dumps.

flush()[source]#

Write all dump buffers out to disk.

setup_hooks(model)[source]#

Define hooking functions on the given torch.nn.Module, but don’t install them.

Parameters

model – torch.nn.Module that serves as the root for recursive names

previous

modelzoo.common.pytorch.dump_context

next

modelzoo.common.pytorch.half_dtype