Training Types

Training Types#

Public API for training package.

CheckpointMetadata#

final class orbax.checkpoint.experimental.v1.training.CheckpointMetadata(step, path, *, metadata, init_timestamp_nsecs=None, commit_timestamp_nsecs=None, custom_metadata=None, metrics=None)[source][source]#

Represents metadata for a single checkpoint (corresponding to a step).

Like its parent, the class has a metadata attribute that is a generic type. The .metadata attribute contains checkpointable-specific metadata. If a PyTree was saved, it will contain PyTreeMetadata, otherwise if `Checkpointable`s were saved, it will be a dictionary mapping names to metadata.

The Orbax checkpointing API provides two symmetric levels of interaction:

  1. Higher level (sequence-of-steps API): Accessed via Checkpointer.

  2. Lower level (individual path API): Accessed via free functions.

CheckpointMetadata objects are returned by both API levels using the same core methods (metadata() and checkpointables_metadata()), reflecting this inherent symmetry.

See superclass documentation for more information, and for a list of base attributes. This class defines several additional attributes that are relevant to checkpoints in a sequence, but not necessarily to a singular checkpoint in isolation.

Example Usage:

from orbax.checkpoint import v1 as ocp

# Higher level (sequence-of-steps API)
with ocp.training.Checkpointer('/path/to/my/checkpoints') as ckptr:
  ckpt_meta = ckptr.metadata(100)

# Lower level (individual path API)
ckpt_meta = ocp.metadata('/path/to/my/checkpoints/100')

# Inspect checkpoint-level properties
print(f'Init time (ns): {ckpt_meta.init_timestamp_nsecs}')
print(f'Commit time (ns): {ckpt_meta.commit_timestamp_nsecs}')
print(f'Custom metadata: {ckpt_meta.custom_metadata}')

# The `.metadata` field contains checkpointable-specific metadata,
# which will be `PyTreeMetadata` or dict[str, CheckpointableMetadata]
# depending on what was saved.
print(f'Checkpointable metadata: {ckpt_meta.metadata}')

See also RootMetadata.

See the parent class, CheckpointMetadata, for base attributes.

Additional Attributes:

step: The step number of the checkpoint. metrics: An optional dictionary containing user-provided metrics saved

alongside the checkpoint.

__init__(step, path, *, metadata, init_timestamp_nsecs=None, commit_timestamp_nsecs=None, custom_metadata=None, metrics=None)[source][source]#

RootMetadata#

final class orbax.checkpoint.experimental.v1.training.RootMetadata(*, directory, custom_metadata=None)[source][source]#

Metadata of a sequence of checkpoint at root level (contains all steps).

This class represents the top-level metadata for an entire checkpointing directory, distinct from step-specific metadata. It associates the physical storage location of the sequence with arbitrary, user-defined information that applies to all steps (e.g., experiment configuration).

Example Usage:

RootMetadata objects are returned by root_metadata().

It can be used to inspect checkpoint-wide information, such as experiment configuration:

import orbax.checkpoint.v1 as ocp
ckptr = ocp.training.Checkpointer('/path/to/my/checkpoints')
root_meta = ckptr.root_metadata()

print(f'Directory: {root_meta.directory}')
print(f'Custom metadata: {root_meta.custom_metadata}')

See also CheckpointMetadata.

directory#

The directory of the checkpoint sequence.

Type:

path_types.Path

custom_metadata#

User-provided custom metadata. An arbitrary JSON-serializable dictionary the user can use to store additional information. The field is treated as opaque by Orbax.

Type:

tree_types.JsonType | None

__delattr__(name)#

Implement delattr(self, name).

__eq__(other)#

Return self==value.

__hash__()#

Return hash(self).

__init__(*, directory, custom_metadata=None)#
__setattr__(name, value)#

Implement setattr(self, name, value).

Contents