tf.keras.distribution.DeviceMesh
Кластер вычислительных устройств для распределённых вычислений.
tf.keras.distribution.DeviceMesh(
shape, axis_names, devices=None
)
Этот API согласован с jax.sharding.Mesh и tf.dtensor.Mesh, которые представляют вычислительные устройства в глобальном контексте.
Дополнительные сведения см. в jax.sharding.Mesh и tf.dtensor.Mesh.
| Аргументы | |
|---|---|
shape | кортеж списков целых чисел. Форма общего DeviceMesh, например, (8,) для распределения только по данным или (4, 2) для распределения по модели + данным. |
axis_names | Список строк. Логическое имя каждой оси для DeviceMesh. Длина axis_names должна соответствовать рангу shape. axis_names будет использоваться для сопоставления/создания TensorLayout при распределении данных и переменных. |
devices | Необязательный список устройств. По умолчанию включает все доступные устройства локально из keras.distribution.list_devices(). |
| Атрибуты | |
|---|---|
axis_names | |
devices | |
shape | |
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/keras/distribution/DeviceMesh
