Preservation Policy

Defines policies for when a checkpoint is preserved.

PreservationContext

class orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationContext

Additional properties for making a save decision.

__eq__(other)

Return self==value.

__hash__ = None

__init__()

PreservationPolicy

class orbax.checkpoint.checkpoint_managers.preservation_policy.PreservationPolicy(*args, **kwargs)

A policy that defines when checkpoints should be preserved.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__init__(*args, **kwargs)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

PreserveAll

class orbax.checkpoint.checkpoint_managers.preservation_policy.PreserveAll

Preserves all checkpoints.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__()

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

LatestN

class orbax.checkpoint.checkpoint_managers.preservation_policy.LatestN(n=None)

Preserves the last n checkpoints. Preserves all checkpoint if n is None.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(n=None)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

EveryNSeconds

class orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSeconds(interval_secs)

Ensures checkpoints are preserved at least after the time interval.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(interval_secs)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

EveryNSteps

class orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNSteps(interval_steps, exact_interval=True, max_to_keep=None)

Preserves checkpoints after at least N steps.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(interval_steps, exact_interval=True, max_to_keep=None)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

CustomSteps

class orbax.checkpoint.checkpoint_managers.preservation_policy.CustomSteps(steps)

Preserves checkpoints at the given steps.

__post_init__(steps_init)

Initializes the internal set of steps after the object is created.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(steps)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

AnyPreservationPolicy

class orbax.checkpoint.checkpoint_managers.preservation_policy.AnyPreservationPolicy(policies)

Applies multiple preservation policies and preserves if any policy preserves.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(policies)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

BestN

class orbax.checkpoint.checkpoint_managers.preservation_policy.BestN(*, get_metric_fn, reverse=False, n=None, keep_checkpoints_without_metrics=True)

A policy that preserves the best checkpoints based on a best_fn.

get_metric_fn:

A function that accepts a nested tree of metrics and returns a scalar value representing the value used for ranking checkpoints.

reverse:

If False (default), checkpoints are sorted in ascending order, according to the best_fn. If True, checkpoints are sorted in descending order. Same as the semantics of built-in sorted() function.

n:

The number of checkpoints to preserve. If None, all checkpoints are preserved. If 0, no checkpoints are preserved.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(*, get_metric_fn, reverse=False, n=None, keep_checkpoints_without_metrics=True)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

LatestDuration

class orbax.checkpoint.checkpoint_managers.preservation_policy.LatestDuration(duration)

Preserves checkpoints that are newer than the given duration.

E.g. retain checkpoints within the last 24 hours:

import datetime
LatestDuration(datetime.timedelta(hours=24))

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(duration)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

EveryNStepsClosest

class orbax.checkpoint.checkpoint_managers.preservation_policy.EveryNStepsClosest(interval_steps, max_to_keep=None)

Preserves checkpoints at steps closest to absolute multiples of N.

This policy maps each checkpoint to its closest nominal target step on a grid defined by interval_steps (i.e. k * interval_steps). For each nominal target, the closest available checkpoint is preserved.

This avoids the error accumulation/drift that can occur with EveryNSteps(exact_interval=False) when checkpoints are irregular.

The last checkpoint is always preserved for final model state and efficient recovery.

should_preserve(checkpoints, *, context)

Indicates which checkpoints should be preserved..

Return type:

Sequence[bool]

__eq__(other)

Return self==value.

__hash__ = None

__init__(interval_steps, max_to_keep=None)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).