grain.checkpoint module

grain.checkpoint module#

APIs for saving and restoring pipeline state.

List of Members#

class grain.checkpoint.CheckpointHandler#

Orbax CheckpointHandler for PyGrain iterators.

restore(directory, item=None, args=None)#

Restores the given iterator from the checkpoint in directory.

Parameters:

  • directory (Path)

  • item (IteratorType | None)

  • args (Any)

Return type:

IteratorType

save(directory, item=None, args=None)#

Saves the given iterator to the checkpoint in directory.

Parameters:

  • directory (Path)

  • item (IteratorType | None)

  • args (Any)

class grain.checkpoint.CheckpointSave(item: Any)#
Parameters:

item (Any)

class grain.checkpoint.CheckpointRestore(item: Any)#
Parameters:

item (Any)

Contents