tf.contrib.checkpoint.split_dependency

View source on GitHub

Creates multiple dependencies with a synchronized save/restore.

tf.contrib.checkpoint.split_dependency(
    component_names,
    component_dtypes,
    fill_save_buffer_fn,
    consume_restore_buffer_fn,
    device
)

Useful when a single op produces Tensors which should each be saved under different objects, or when Tensors saved with many different objects need to be restored together as inputs to a single op (i.e. an object which uses a single fused op may be swapped out for a subgraph of objects, and these two programs are checkpoint compatible).

Args:

  • component_names: A sequence of names for the split dependencies. fill_save_buffer_fn must add these keys to the dictionary it is passed, and consume_restore_buffer_fn will receive a dictionary with these keys.
  • component_dtypes: Data types for the Tensors being saved and restored, a sequence corresponding to component_names.
  • fill_save_buffer_fn: A function which takes an empty dictionary as an argument and adds Tensors with component_names as keys. These Tensors will be saved as if they were individual variables.
  • consume_restore_buffer_fn: A function which takes a dictionary with component_names as keys mapping to restored individual Tensors and returns a restore op (or if executing eagerly, runs the restoration and may return None).
  • device: The device on which to run save and restore operations.

Returns:

A dictionary mapping from names to Trackable objects. If one is reachable from an object as a dependency, the others should be too; adding dependencies on some but not all of the objects will result in errors.