tf.distribute.experimental.partitioners.MinSizePartitioner

Partitioner that allocates a minimum size per shard.

Inherits From: Partitioner

Used in the notebooks

Used in the tutorials

This partitioner ensures each shard has at least min_shard_bytes, and tries to allocate as many shards as possible, i.e., keeping shard size as small as possible. The maximum number of such shards (upper bound) is given by max_shards.

Examples:

partitioner = MinSizePartitioner(min_shard_bytes=4, max_shards=2)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[2, 1]
partitioner = MinSizePartitioner(min_shard_bytes=4, max_shards=10)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[6, 1]

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

min_shard_bytes Minimum bytes of each shard. Defaults to 256K.
max_shards Upper bound on the number of shards. Defaults to 1.
bytes_per_string If the partition value is of type string, this provides an estimate of how large each string is.

Methods

__call__

View source