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.

byte_size_load_fn(op)[source]¶

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.