tf.distribute.experimental.partitioners.FixedShardsPartitioner
Partitioner that allocates a fixed number of shards.
Inherits From: Partitioner
tf.distribute.experimental.partitioners.FixedShardsPartitioner(
num_shards
)
Examples:
# standalone usage:
partitioner = FixedShardsPartitioner(num_shards=2)
partitions = partitioner(tf.TensorShape([10, 3]), tf.float32)
[2, 1]
# use in ParameterServerStrategy
# strategy = tf.distribute.experimental.ParameterServerStrategy(
# cluster_resolver=cluster_resolver, variable_partitioner=partitioner)
| Args | |
|---|---|
num_shards |
int, number of shards to partition. |
Methods
__call__
__call__(
shape, dtype, axis=0
)
Partitions the given shape and returns the partition results.
Examples of a partitioner that allocates a fixed number of shards:
partitioner = FixedShardsPartitioner(num_shards=2)
partitions = partitioner(tf.TensorShape([10, 3], tf.float32), axis=0)
print(partitions) # [2, 0]
| Args | |
|---|---|
shape |
a tf.TensorShape, the shape to partition. |
dtype |
a tf.dtypes.Dtype indicating the type of the partition value. |
axis |
The axis to partition along. Default: outermost axis. |
| Returns | |
|---|---|
| A list of integers representing the number of partitions on each axis, where i-th value correponds to i-th axis. |
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/distribute/experimental/partitioners/FixedShardsPartitioner