Training step timing¶

Opt-in, per-stage timing of the training step.

Timing is off unless the METATRAIN_TIMING environment variable is set (or enable() is called), in which case every timed() block accumulates into a process-global table that report() formats. It exists to answer one question before the data pipeline is optimized: which fraction of a training step is spent waiting for the input pipeline rather than running the model.

timed_iter() splits the training loop into the two top-level stages loader (waiting for the next batch) and step (everything the loop body does with it); the stages in STEP_BREAKDOWN break step down further, and are therefore reported as a fraction of loader + step.

The COLLATE_STAGES are recorded inside DataLoader workers, whose accumulators die with the worker process; they are only reported when running with num_workers=0. With workers, that time shows up as loader wait in the main process instead. Each collate transform is timed separately under TRANSFORM_PREFIX, which is what says whether the expensive part of collation is the neighbor lists, the augmentation, or something else.

metatrain.utils.timing.TOP_STAGES = ('loader', 'step')¶: The two halves of a training step: waiting for data, and using it.

metatrain.utils.timing.STEP_BREAKDOWN = ('unpack', 'h2d', 'forward', 'loss', 'backward', 'optimizer')¶: Stages the step stage is made of, in the order they run.

metatrain.utils.timing.COLLATE_STAGES = ('group_and_join', 'transforms', 'serialize')¶: Stages inside the collate function, i.e. inside the loader wait.

metatrain.utils.timing.TRANSFORM_PREFIX = 'transforms/'¶: Prefix of the per-transform stages that break transforms down.

metatrain.utils.timing.enable() → None[source]¶

Turn timing on for this process, as METATRAIN_TIMING does.

Useful for scripts that want timings without having to set the environment variable before importing metatrain.

Return type:: None

metatrain.utils.timing.reset() → None[source]¶

Forget all recorded timings and counters.

Return type:: None

metatrain.utils.timing.timed(stage: str) → ContextManager[None, bool | None][source]¶

Time the enclosed block, or do nothing at all when timing is off.

Parameters:: stage (str) – Name of the stage to accumulate into.
Returns:: A context manager around the timed block.
Return type:: ContextManager[None, bool | None]

metatrain.utils.timing.timed_transform(transform: Any) → ContextManager[None, bool | None][source]¶

Time one collate transform, in a stage named after the callable.

The name is only built when timing is on, so the disabled path stays free.

Parameters:: transform (Any) – The collate transform about to run.
Returns:: A context manager around the timed block.
Return type:: ContextManager[None, bool | None]

metatrain.utils.timing.timed_iter(iterable: Iterable[Any], stage: str, body: str) → Iterator[Any][source]¶

Iterate iterable, attributing waits to stage and the loop body to body.

Parameters:

iterable (Iterable[Any]) – The iterable to consume, e.g. a DataLoader.
stage (str) – Stage name for the wait before each item.
body (str) – Stage name for what the consumer does with each item.

Returns:

An iterator over iterable.

Return type:

Iterator[Any]

metatrain.utils.timing.count_systems(systems: List[System]) → None[source]¶