Variable Partitioners for Sharding

tf.fixed_size_partitioner(num_shards, axis=0)

Partitioner to specify a fixed number of shards along given axis.

Args:
  • num_shards: int, number of shards to partition variable.
  • axis: int, axis to partition on.
Returns:

A partition function usable as the partitioner argument to variable_scope, get_variable, and get_partitioned_variable_list.


tf.variable_axis_size_partitioner(max_shard_bytes, axis=0, bytes_per_string_element=16, max_shards=None)

Get a partitioner for VariableScope to keep shards below max_shard_bytes.

This partitioner will shard a Variable along one axis, attempting to keep the maximum shard size below max_shard_bytes. In practice, this is not always possible when sharding along only one axis. When this happens, this axis is sharded as much as possible (i.e., every dimension becomes a separate shard).

If the partitioner hits the max_shards limit, then each shard may end up larger than max_shard_bytes. By default max_shards equals None and no limit on the number of shards is enforced.

One reasonable value for max_shard_bytes is (64 << 20) - 1, or almost 64MB, to keep below the protobuf byte limit.

Args:
  • max_shard_bytes: The maximum size any given shard is allowed to be.
  • axis: The axis to partition along. Default: outermost axis.
  • bytes_per_string_element: If the Variable is of type string, this provides an estimate of how large each scalar in the Variable is.
  • max_shards: The maximum number of shards in int created taking precedence over max_shard_bytes.
Returns:

A partition function usable as the partitioner argument to variable_scope, get_variable, and get_partitioned_variable_list.

Raises:
  • ValueError: If any of the byte counts are non-positive.

tf.min_max_variable_partitioner(max_partitions=1, axis=0, min_slice_size=262144, bytes_per_string_element=16)

Partitioner to allocate minimum size per slice.

Returns a partitioner that partitions the variable of given shape and dtype such that each partition has a minimum of min_slice_size slice of the variable. The maximum number of such partitions (upper bound) is given by max_partitions.

Args:
  • max_partitions: Upper bound on the number of partitions. Defaults to 1.
  • axis: Axis along which to partition the variable. Defaults to 0.
  • min_slice_size: Minimum size of the variable slice per partition. Defaults to 256K.
  • bytes_per_string_element: If the Variable is of type string, this provides an estimate of how large each scalar in the Variable is.
Returns:

A partition function usable as the partitioner argument to variable_scope, get_variable, and get_partitioned_variable_list.