tf.keras.distribution.DataParallel

Distribution for data parallelism.

tf.keras.distribution.DataParallel(
    device_mesh=None, devices=None, auto_shard_dataset=True
)

You can choose to create this instance by either specifying the device_mesh or devices arguments (but not both).

The device_mesh argument is expected to be a DeviceMesh instance, and is expected to be 1D only. In case that the mesh has multiple axes, then the first axis will be treated as the data parallel dimension (and a warning will be raised).

When a list of devices are provided, they will be used to construct a 1D mesh.

When both mesh and devices are absent, then list_devices() will be used to detect any available devices and create a 1D mesh from them.

device_mesh Optional DeviceMesh instance.
devices Optional list of devices.
auto_shard_dataset Automatically shard the dataset amongst processes in a multi-process setting. Set to False if the dataset is already sharded across hosts. Defaults to True.

auto_shard_dataset

batch_dim_name

data_shard_id ID of the data shard for the current process.
device_mesh

num_model_replicas Number of model replicas.
num_processes Total number of processes in the cluster.

Methods

distribute_dataset

View source

distribute_dataset(
    dataset
)

Create a distributed dataset from the original global dataset.

Args
dataset the original global dataset instance.

Returns
If auto_shard_dataset is True, returns a sharded dataset that only produces data for the current local worker/process. Otherwise, returns the original dataset.

Raises
ValueError if auto-sharding is requested in a multi-process setting, but the dataset type is not supported.

distribute_tf_dataset

View source

distribute_tf_dataset(
    dataset
)

get_data_layout

View source

get_data_layout(
    data_shape
)

Retrieve the TensorLayout for the input data.

Args
data_shape shape for the input data in list or tuple format.

Returns
The TensorLayout for the data, which can be used by backend.distribute_value() to redistribute a input data.

get_tensor_layout

View source

get_tensor_layout(
    path
)

Retrieve the TensorLayout for the intermediate tensor.

Args
path a string path for the corresponding tensor.

return: The TensorLayout for the intermediate tensor, which can be used by backend.relayout() to reshard the tensor. Could also return None.

get_variable_layout

View source

get_variable_layout(
    variable
)

Retrieve the TensorLayout for the variable.

Args
variable A Variable instance.

return: The TensorLayout for the variable, which can be used by backend.distribute_value() to redistribute a variable.

scope

View source

@contextlib.contextmanager
scope()

Context manager to make the Distribution current.