Uneven Partition Ps Strategy

Partitioned PS StrategyBuilder with Greedy Load Balancer.

class UnevenPartitionedPS(local_proxy_variable=False, sync=True, staleness=0)

Bases: autodist.strategy.base.StrategyBuilder

Partitioned PS StrategyBuilder with Greedy Load Balancer.

Determine the number of partitions for each partition-able variable by finding its minimum divisor along the first axis.

Then, use a greedy load balancer (determined by memory usage) to assign the partitions to Parameter Servers. This means that, unlike the standard PS StrategyBuilder, that a variable can be spread across multiple servers using this StrategyBuilder.

build(graph_item, resource_spec)

Generate the Strategy.

static get_num_shards(var)

Gets the minimum number of shards for a variable.

byte_size_load_fn(op)

Load function that computes the byte size of a single-output Operation.

Copied (with modifications) from tensorflow.contrib.training.python.training.device_setter.

This is intended to be used with “Variable” ops, which have a single Tensor output with the contents of the variable. However, it can also be used for calculating the size of any op that has a single output.

Intended to be used with GreedyLoadBalancingStrategy.

Parameters

op – An Operation with a single output, typically a “Variable” op.

Returns

The number of bytes in the output Tensor.

Raises

ValueError – if op does not have a single output, or if the shape of the single output is not fully-defined.