tf.tpu.experimental.embedding.TPUEmbeddingV2

Среднеуровневый API TPUEmbedding, работающий на TPU с ускорителем разреженных вычислений.

Просмотр псевдонимов

Псевдонимы совместимости для миграции

Дополнительные сведения см. в руководстве по миграции.

tf.compat.v1.tpu.experimental.embedding.TPUEmbeddingV2

tf.tpu.experimental.embedding.TPUEmbeddingV2(
    feature_config: Union[tf.tpu.experimental.embedding.FeatureConfig, Iterable],
    optimizer: Optional[tpu_embedding_v2_utils._Optimizer] = None,
    pipeline_execution_with_tensor_core: bool = False
)

Аргументы
`feature_config`	Вложенная структура конфигураций `tf.tpu.experimental.embedding.FeatureConfig`.
`optimizer`	Экземпляр одного из оптимизаторов `tf.tpu.experimental.embedding.SGD`, `tf.tpu.experimental.embedding.Adagrad` или `tf.tpu.experimental.embedding.Adam`. При создании вне TPUStrategy может быть установлен в None для предотвращения создания переменных оптимизатора, что полезно для оптимизации потребления памяти при экспорте модели для использования, где переменные слотов не нужны.
`pipeline_execution_with_tensor_core`	Если True, вычисления TPU embedding будут перекрываться с вычислениями TensorCore (и, следовательно, будут иметь отставание на один шаг). Установите в True для повышения производительности.

Исключения
`ValueError`	Если оптимизатор не является одним из tf.tpu.experimental.embedding.(SGD, Adam или Adagrad) или None при создании внутри TPUStrategy.
`RuntimeError`	Если не создано внутри TPUStrategy.

Атрибуты
`embedding_table_shards`	Возвращает словарь таблиц embedding, индексированный по `TableConfig`.
`embedding_tables`	Возвращает словарь таблиц embedding, индексированный по `TableConfig`.
`variables`	Возвращает словарь переменных, индексированный по `TableConfig`, а затем по имени слота.

Методы

`apply_gradients`

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

apply_gradients(
    gradients: Any, preserved_outputs: Dict[str, PartitionedCsrFormatTensor]
)

Применяет обновление градиента к таблицам embedding.

Если градиент None передан в любой позиции вложенной структуры, то для этой функции применяется обновление градиента с нулевым градиентом. Для оптимизаторов, таких как SGD или Adagrad, это равносильно тому, что обновление вообще не применяется. Для ленивого Adam и других оптимизаторов, применяемых разреженно, с затуханием, убедитесь, что вы понимаете эффект применения нулевого градиента.

Аргументы
`gradients`	Вложенная структура градиентов, структура которой соответствует структуре `feature_config`, переданной в этот объект.
`preserved_outputs`	Словарь PartitionedCsrFormatTensor, полученный со второго выхода вызова embedding lookup.

Исключения
`RuntimeError`	если не построено.
`ValueError`	Если передан не-`tf.Tensor` не-`None` градиент или передан `tf.Tensor` неправильной формы. Также, если размер любого набора в `gradients` не соответствует соответствующему набору в `feature_config`.
`TypeError`	Если тип любого набора в `gradients` не соответствует соответствующему набору в `feature_config`.

`build`

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

build()

Создает переменные и переменные слотов для TPU embedding.

`dequeue`

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

dequeue(
    partitioned_tensors: Tuple[Dict[str, PartitionedCsrFormatTensor], int, int]
) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]

Выполняет поиск по embedding.

`embedding_lookup`

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

embedding_lookup(
    features: Any, weights: Optional[Any] = None
) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]

Выполняет поиск по входному признаку.

Аргументы

features Вложенная структура tf.Tensors, tf.SparseTensors или tf.RaggedTensors, с той же структурой, что и feature_config. Входные данные будут приведены к типу tf.int32. Поддерживается только один тип из tf.SparseTensor или tf.RaggedTensor за один вызов.

weights Если не None, вложенная структура tf.Tensors, tf.SparseTensors или tf.RaggedTensors, соответствующая вышеуказанной, за исключением того, что тензоры должны быть плавающего типа (и они будут приведены к типу tf.float32). Для tf.SparseTensor мы предполагаем, что indices одинаковы для параллельных записей из features, и аналогично для tf.RaggedTensor мы предполагаем, что row_splits одинаковы.

Аргументы
`features`	Вложенная структура `tf.Tensor`s, `tf.SparseTensor`s или `tf.RaggedTensor`s, с той же структурой, что и `feature_config`. Входные данные будут приведены к типу `tf.int32`. Поддерживается только один тип из `tf.SparseTensor` или `tf.RaggedTensor` за один вызов.
`weights`	Если не `None`, вложенная структура `tf.Tensor`s, `tf.SparseTensor`s или `tf.RaggedTensor`s, соответствующая вышеуказанной, за исключением того, что тензоры должны быть плавающего типа (и они будут приведены к типу `tf.float32`). Для `tf.SparseTensor` мы предполагаем, что `indices` одинаковы для параллельных записей из `features`, и аналогично для `tf.RaggedTensor` мы предполагаем, что row_splits одинаковы.

Исключения
`ValueError`	Если входной признак не является типом Tensor, SparseTensor или RaggedTensor.
`TypeError`	Если тип любого набора в `features` не соответствует соответствующему набору в `feature_config`. Аналогично для `weights`, если не `None`.

Возвращаемые значения
`packed_activations`	Результаты поиска по embedding, упакованные в том же порядке, что и входной признак.
`packed_output`	Словарь PartitionedCsrFormatTensors.

`enqueue`

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

enqueue(
    features: Any, weights: Optional[Any] = None, device: Optional[str] = None
) -> Any

Предварительная обработка признаков на хосте.

`preprocess_features`

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

@staticmethod
preprocess_features(
    num_replicas_in_sync: int,
    max_ids_per_chip_per_sample: int,
    max_minibatches_per_sc: int,
    num_sc_per_chip: int,
    num_sc_shards: int,
    stacked_table_to_tables: Dict[str, Any],
    table_to_stacked_table_offset: Dict[str, Tuple[str, int, int]],
    table_to_sample_count: Dict[str, int],
    feature_to_sample_offset: Dict[str, int],
    flat_features: Any,
    flat_inputs: Any,
    flat_weights: Optional[Any] = None
) -> Any

Функция для предварительной обработки признаков.

`call`

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

__call__(
    features: Any, weights: Optional[Any] = None
) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]

Вызов среднеуровневого API для выполнения поиска по embedding.

© 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/tpu/experimental/embedding/TPUEmbeddingV2