tf.distribute.experimental.partitioners.MaxSizePartitioner

Разделитель, который поддерживает фрагменты ниже max_shard_bytes.

Наследуется от: Partitioner

tf.distribute.experimental.partitioners.MaxSizePartitioner(
    max_shard_bytes, max_shards=None, bytes_per_string=16
)

Этот разделитель гарантирует, что каждый фрагмент имеет не более max_shard_bytes, и пытается выделить как можно меньше фрагментов, то есть, сохраняя размер фрагмента как можно большим.

Если разделитель достигает предела max_shards, каждый фрагмент может оказаться больше, чем max_shard_bytes. По умолчанию max_shards равно None и ограничение на количество фрагментов не применяется.

Примеры:

partitioner = MaxSizePartitioner(max_shard_bytes=4)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[6, 1]
partitioner = MaxSizePartitioner(max_shard_bytes=4, max_shards=2)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[2, 1]
partitioner = MaxSizePartitioner(max_shard_bytes=1024)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[1, 1]

# use in ParameterServerStrategy
# strategy = tf.distribute.experimental.ParameterServerStrategy(
#   cluster_resolver=cluster_resolver, variable_partitioner=partitioner)

Методы

__call__

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

__call__(
    shape, dtype, axis=0
)

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

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

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