View source on GitHub |
Saves given named tensor slices in a sharded, multi-client safe fashion.
tf.experimental.dtensor.sharded_save(
mesh: tf.experimental.dtensor.Mesh
,
file_prefix: Union[str, tf.Tensor
],
tensor_names: Union[List[str], tf.Tensor
],
shape_and_slices: Union[List[str], tf.Tensor
],
tensors: List[Union[ops.Tensor, tf_variables.Variable]]
)
The method makes sure the checkpoint directory state is correct in a sharded mutli-client saving. Namely, we place a barrier after SaveV2 to make sure every client has done writing the files. And another one after MergeV2Checkpoints to make sure all Metadata is properly merged.
Upon existing, the checkpoint is completed and the all directory operations are done.
Returns | |
---|---|
A MergeV2Checkpoints op that merged all Metadata. |