AllToAll

classe final pública AllToAll

Um Op para trocar dados entre réplicas de TPU.

Em cada réplica, a entrada é dividida em blocos `split_count` ao longo de `split_dimension` e enviada para as outras réplicas com group_assignment. Depois de receber `split_count` - 1 blocos de outras réplicas, concatenamos os blocos ao longo de `concat_dimension` como saída.

Por exemplo, suponha que haja 2 réplicas de TPU: a réplica 0 recebe a entrada: `[[A, B]]` a réplica 1 recebe a entrada: `[[C, D]]`

group_assignment=`[[0, 1]]` concat_dimension=0 split_dimension=1 split_count=2

saída da réplica 0: `[[A], [C]]` saída da réplica 1: `[[B], [D]]`

Métodos públicos

Saída <T>
comoSaída ()
Retorna o identificador simbólico de um tensor.
estático <T> TodosParaTodos <T>
create ( Escopo do escopo, Operando <T> entrada, Operando <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)
Método de fábrica para criar uma classe que envolve uma nova operação AllToAll.
Saída <T>
saída ()
O resultado trocado.

Métodos Herdados

Métodos públicos

Public Output <T> asOutput ()

Retorna o identificador simbólico de um tensor.

As entradas para operações do TensorFlow são saídas de outra operação do TensorFlow. Este método é usado para obter um identificador simbólico que representa o cálculo da entrada.

public static AllToAll <T> create ( Escopo do escopo, Operando <T> entrada, Operando <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)

Método de fábrica para criar uma classe que envolve uma nova operação AllToAll.

Parâmetros
alcance escopo atual
entrada A entrada local para a soma.
tarefa de grupo Um tensor int32 com forma [num_groups, num_replicas_per_group]. `group_assignment[i]` representa os ids de réplica no º subgrupo.
concatDimension O número da dimensão a ser concatenado.
splitDimension O número da dimensão a ser dividido.
splitCount O número de divisões, este número deve ser igual ao tamanho do subgrupo(group_assignment.get_shape()[1])
Devoluções
  • uma nova instância de AllToAll

saída pública <T> saída ()

O resultado trocado.