View source on GitHub |
Restores from checkpoint_prefix to name based DTensors.
tf.experimental.dtensor.name_based_restore(
mesh: tf.experimental.dtensor.Mesh
,
checkpoint_prefix: str,
name_tensor_dict: Dict[str, Union[tensor_lib.Tensor, tf_variables.Variable]]
)
It is required to have already-initialized DTensor variables that have same shape/dtype for the tensors being restored.
Also, we currently only support a named based restore on a single mesh.
Returns | |
---|---|
A dictionary of name to its restored DTensor value. |