tf.experimental.dtensor.Mesh
Представляет конфигурацию Mesh над определённым списком измерений Mesh.
tf.experimental.dtensor.Mesh(
dim_names: List[str],
global_device_ids: np.ndarray,
local_device_ids: List[int],
local_devices: List[tf.compat.v1.DeviceSpec],
mesh_name: str = '',
global_devices: Optional[List[tf_device.DeviceSpec]] = None
)
Массив состоит из именованных измерений с размерами, которые описывают, как набор устройств расположен. Определение раскладов тензоров в терминах измерений массива позволяет нам эффективно определять коммуникацию, необходимую при вычислении операции с тензорами разных раскладов.
Массив предоставляет информацию не только о размещении тензоров, но и о топологии подлежащих устройств. Например, мы можем сгруппировать 8 TPU в 1-мерный массив для параллелизма данных или в 2x4 сетку для (двухстороннего) параллелизма данных и (четырёхстороннего) параллелизма модели.
Примечание: утилитыdtensor.create_meshиdtensor.create_distributed_meshпредоставляют более простой API для создания массивов для одно- или многоклиентских случаев.
| Аргументы | |
|---|---|
dim_names | Список строк, указывающих имена измерений. |
global_device_ids | Массив ndarray глобальных идентификаторов устройств используется для составления DeviceSpecs, описывающих массив. Форма этого массива определяет размер каждого измерения массива. Значения в этом массиве должны последовательно увеличиваться от 0. Этот аргумент одинаков для каждого клиента DTensor. |
local_device_ids | Список локальных идентификаторов устройств, равный подмножеству значений в global_device_ids. Они указывают положение локальных устройств в глобальном массиве. Разные клиенты DTensor должны содержать различные локальные идентификаторы устройств. Все local_device_ids от всех клиентов DTensor должны покрывать каждый элемент в global_device_ids. |
local_devices | Список устройств, размещённых локально. Элементы соответствуют 1:1 элементам local_device_ids. |
mesh_name | Имя массива. В настоящее время это редко используется и в основном используется для указания, является ли это массив на основе CPU, GPU или TPU. global_devices (необязательно): Список глобальных устройств. Устанавливается при использовании нескольких массивов устройств. |
| Атрибуты | |
|---|---|
dim_names | |
name | |
size | |
Методы
as_proto
as_proto() -> layout_pb2.MeshProto
Возвращает протобуфер массива.
contains_dim
contains_dim(
dim_name: str
) -> bool
Возвращает True, если Mesh содержит заданное имя измерения.
device_type
device_type() -> str
Возвращает тип устройства Mesh.
dim_size
dim_size(
dim_name: str
) -> int
Возвращает размер измерения.
from_proto
@staticmethod
from_proto(
proto: layout_pb2.MeshProto
) -> 'Mesh'
Создаёт экземпляр массива из входного proto.
from_string
@staticmethod
from_string(
mesh_str: str
) -> 'Mesh'
Создаёт экземпляр массива из входного proto.
host_mesh
host_mesh()
Возвращает сопоставленный 1:1 хост-массив.
is_remote
is_remote() -> bool
Возвращает True, если Mesh содержит только удалённые устройства.
local_device_ids
local_device_ids() -> List[int]
Возвращает список локальных идентификаторов устройств.
local_device_locations
local_device_locations() -> List[Dict[str, int]]
Возвращает список локаций локальных устройств.
Расположение устройства — это словарь из имён измерений и индексов на этих измерениях.
local_devices
local_devices() -> List[str]
Возвращает список локальных спецификаций устройств, представленных в виде строк.
min_global_device_id
min_global_device_id() -> int
Возвращает минимальный глобальный идентификатор устройства.
num_local_devices
num_local_devices() -> int
Возвращает количество локальных устройств.
shape
shape() -> List[int]
Возвращает форму массива.
to_string
to_string() -> str
Возвращает строковое представление массива Mesh.
unravel_index
unravel_index()
Возвращает словарь из идентификатора устройства в {имя_измерения: индекс_измерения}.
Например, для массива 3x2, возвращает это:
{ 0: {'x': 0, 'y', 0},
1: {'x': 0, 'y', 1},
2: {'x': 1, 'y', 0},
3: {'x': 1, 'y', 1},
4: {'x': 2, 'y', 0},
5: {'x': 2, 'y', 1} }
__contains__
__contains__(
dim_name: str
) -> bool
__eq__
__eq__(
other
)
Возвращает self==value.
__getitem__
__getitem__(
dim_name: str
) -> MeshDimension
© 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/versions/r2.9/api_docs/python/tf/experimental/dtensor/Mesh
