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
