tff.simulation.FileProgramStateManager

A ProgramStateManager backed by a file system.

Inherits From: ProgramStateManager

The FileProgramStateManager is a utility that can be used to save and load program state to a file system. This implemenation is compatible with any nested structure supported by tf.convert_to_tensor and is intended to support fault tolerance in programs. In particular, it is intended to only restart the same simulation, run with the same version of TensorFlow Federated.

root_dir A path on the filesystem to save program state.
prefix A string to use as the prefix for filenames.
keep_total An integer representing the total number of program states to keep. If the value is zero or smaller, all program states will be kept.
keep_first A boolean indicating if the first program state should be kept, irrespective of whether it is the oldest program state or not. This is desirable in settings where you would like to ensure full reproducibility of the simulation, especially in settings where model weights or optimizer states are initialized randomly. By loading from the initial program state, one can avoid re-initializing and obtaining different results.

Methods

load

View source

Returns the program state for the given version.

Args
version A integer representing the version of a saved program state.

Raises
VersionNotFoundError If there is no program state for the given version.

load_latest

View source

Returns the latest saved program state and version or (None, 0).

Returns
A tuple of the latest saved (program state, version) or (None, 0) if there is no latest saved program state.

save

View source

Saves program_state for the given version.

Args
program_state The program state to save.
version A strictly increasing integer representing the version of a saved program_state

set_structure

View source

Configures a structure to use when loading program state.

The structure must be set before calling load.

Args
structure A nested structure which tf.convert_to_tensor supports to use as a template when calling load.

versions

View source

Returns a list of saved versions or None.

Returns
A list of saved versions or None if there is no saved program state.