Ps Lb Strategy

PS Load Balancing StrategyBuilder.

class PSLoadBalancing(local_proxy_variable=False, sync=True, staleness=0)[source]

Bases: autodist.strategy.base.StrategyBuilder

PS StrategyBuilder with Greedy Load Balancing.

The Load Balancing is determined by total memory usage for storing variables, i.e. we always assign a variable to the current lowest-memory-usage Parameter Server.

build(graph_item, resource_spec)[source]

Generate the Strategy.


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

Copied (with modifications) from

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.


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


The number of bytes in the output Tensor.


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