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
)

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

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

Обратитесь к Концепции DTensor для подробного обсуждения и примеров.

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

[(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

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

contains_dim

contains_dim()

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

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

coords

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

coords(
    device_idx: int
) -> tf.Tensor

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

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

Возвращает тип устройства Mesh.

dim_size

dim_size()

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

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

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) -> Список[int]

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

to_string

to_string()

to_string(self: tensorflow.python._pywrap_dtensor_device.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} }

use_xla_spmd

use_xla_spmd()

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

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

__contains__

__contains__()

contains(self: tensorflow.python._pywrap_dtensor_device.Mesh, dim_name: строка) -> 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