tf.distribute.experimental.partitioners.Partitioner

Базовый класс Partitioner: все partitioner наследуются от этого класса.

Partitioner должен реализовать метод __call__ со следующим сигнатуром:

def __call__(self, shape, dtype, axis=0):
  # Partitions the given `shape` and returns the partition results.
  # See docstring of `__call__` method for the format of partition results.

Методы

__call__

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

__call__(
    shape, dtype, axis=0
)

Разделяет указанное shape и возвращает результаты разделения.

Примеры partitioner, который выделяет фиксированное количество фрагментов:

partitioner = FixedShardsPartitioner(num_shards=2)
partitions = partitioner(tf.TensorShape([10, 3], tf.float32), axis=0)
print(partitions) # [2, 0]