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[Union[tf_device.DeviceSpec, str]],
    mesh_name: str = '',
    global_devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
    use_xla_spmd: bool = USE_XLA_SPMD
)

Mesh состоит из именованных измерений с размерами, которые описывают расположение набора устройств. Определение макетов тензоров через измерения Mesh позволяет эффективно определять необходимые коммуникации при вычислении операций с тензорами, имеющими разные макеты.

Mesh предоставляет информацию не только о размещении тензоров, но и о топологии базовых устройств. Например, можно сгруппировать 8 TPU в одномерный массив для параллелизма по данным или в сетку 2x4 для параллелизма по данным (2 направления) и параллелизма модели (4 направления).

Подробное обсуждение и примеры см. в разделе Основные понятия DTensor.

Примечание: утилиты dtensor.create_mesh и dtensor.create_distributed_mesh предоставляют более простой API для создания Mesh в сценариях с одним или несколькими клиентами.

[(device_id / (b*c*d)) % a,
 (device_id / (c*d))   % b,
 (device_id / (d))     % c,
 (device_id)           % d]

Методы

as_proto

as_proto()

as_proto(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> tensorflow::dtensor::MeshProto

Возвращает сообщение protobuf MeshProto.

contains_dim

contains_dim()

contains_dim(self: tensorflow.python._pywrap_dtensor_device.Mesh, dim_name: str) -> bool

Возвращает True, если Mesh содержит указанное имя измерения.

coords

Посмотреть исходный код

coords(
    device_idx: int
) -> tf.Tensor

Преобразует индекс устройства в тензор координат Mesh.

device_location

device_location()

device_location(self: tensorflow.python._pywrap_dtensor_device.Mesh, arg0: int) -> List[int]

device_type

device_type()

device_type(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> str

Возвращает device_type объекта Mesh.

dim_size

dim_size()

dim_size(self: tensorflow.python._pywrap_dtensor_device.Mesh, dim_name: str) -> int

Возвращает размер измерения Mesh.

from_device

Посмотреть исходный код

@classmethod
from_device(
    device: str
) -> 'Mesh'

Создаёт одноустройственную сетку из строки устройства.

from_proto

Посмотреть исходный код

@classmethod
from_proto(
    proto: layout_pb2.MeshProto
) -> 'Mesh'

Создаёт экземпляр сетки из входного значения proto.

from_string

Посмотреть исходный код

@classmethod
from_string(
    mesh_str: str
) -> 'Mesh'

global_device_ids

Посмотреть исходный код

global_device_ids() -> np.ndarray

Возвращает список глобальных устройств в виде массива.

global_devices

global_devices()

global_devices(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> List[str]

Возвращает список спецификаций глобальных устройств в виде строк.

host_mesh

Посмотреть исходный код

host_mesh() -> 'Mesh'

Возвращает хостовую сетку.

is_remote

is_remote()

is_remote(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> bool

Возвращает True, если сетка содержит только удалённые устройства.

is_single_device

is_single_device()

is_single_device(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> bool

Возвращает True, если сетка представляет нераспределённое устройство.

local_device_ids

local_device_ids()

local_device_ids(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> Span[int]

Возвращает список идентификаторов локальных устройств.

local_device_locations

Посмотреть исходный код

local_device_locations() -> List[Dict[str, int]]

Возвращает список расположений локальных устройств.

Расположение устройства — это словарь, сопоставляющий имена измерений с индексами в этих измерениях.

local_devices

local_devices()

local_devices(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> Span[str]

Возвращает список спецификаций локальных устройств в виде строк.

min_global_device_id

min_global_device_id()

min_global_device_id(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> int

Возвращает минимальный глобальный идентификатор устройства.

num_local_devices

num_local_devices()

num_local_devices(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> int

Возвращает количество локальных устройств.

shape

shape()

shape(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> List[int]

Возвращает форму сетки.

to_string

to_string()

to_string(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> str

Возвращает строковое представление сетки.

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} }

use_xla_spmd

use_xla_spmd()

use_xla_spmd(self: tensorflow.python._pywrap_dtensor_device.Mesh) -> bool

Возвращает True, если сетка будет использовать XLA для SPMD вместо SPMD в DTensor.

__contains__

__contains__()

contains(self: tensorflow.python._pywrap_dtensor_device.Mesh, dim_name: str) -> bool

__eq__

__eq__()

eq(self: tensorflow.python._pywrap_dtensor_device.Mesh, arg0: tensorflow.python._pywrap_dtensor_device.Mesh) -> bool

__getitem__

Посмотреть исходный код

__getitem__(
    dim_name: str
) -> MeshDimension