tf.compat.v1.tpu.batch_parallel
Разделяет computation по размеру пакетной размерности для параллельного выполнения.
tf.compat.v1.tpu.batch_parallel(
computation: Callable[..., Any],
inputs: Optional[List[List[Optional[core_types.Tensor]]]] = None,
num_shards: int = 1,
infeed_queue: Optional[tpu_feed.InfeedQueue] = None,
device_assignment: Optional[tf.tpu.experimental.DeviceAssignment] = None,
name: Optional[Text] = None,
xla_options: Optional[tf.tpu.XLAOptions] = None
)
Обёртка для удобства работы со shard().
inputs должен быть списком тензоров или None (эквивалентно пустому списку). Каждый вход разбивается на num_shards части по 0-й размерности, и вычисление применяется к каждой части параллельно.
Тензоры транслируются на все части, если они лексически заключены в computation. Например:
x = tf.constant(7) def computation(): return x + 3 ... = shard(computation, ...)
Выходы всех частей конкатенируются обратно по их 0-й размерности.
Входы и выходы вычисления должны быть тензорами с размером хотя бы 1.
| Аргументы | |
|---|---|
computation | Функция Python, которая строит вычисление, которое нужно применить к каждой части входных данных. |
inputs | Список входных тензоров или None (эквивалентно пустому списку). 0-й размер каждого тензора должен быть кратен num_shards. |
num_shards | Количество частей. |
infeed_queue | Если не None, то кортеж аргументов, которые нужно добавить в качестве входов в computation. |
device_assignment | Если не None, то объект DeviceAssignment, описывающий сопоставление между логическими ядрами в вычислении и физическими ядрами в топологии TPU. Использует значение по умолчанию для распределения устройств, если None. DeviceAssignment можно опустить, если каждая часть вычисления использует только одно ядро, и либо есть только одна часть, либо количество частей равно количеству ядер в системе TPU. |
name | (Устарело) Ничего не делает. |
xla_options | Экземпляр tpu.XLAOptions, который указывает параметры, передаваемые компилятору XLA. Используйте None для параметров по умолчанию. |
| Возвращаемые значения | |
|---|---|
| Список выходных тензоров. |
| Исключения | |
|---|---|
ValueError | Если num_shards <= 0 |
© 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/compat/v1/tpu/batch_parallel