tensorpack.graph_builder package

These are some useful functions if you need to write your own trainers. Otherwise you probably don’t need to use them.

class tensorpack.graph_builder.DataParallelBuilder(towers)[source]

Bases: tensorpack.graph_builder.training.GraphBuilder

__init__(towers)[source]
Parameters:towers (list[int]) – list of GPU ids.
static call_for_each_tower(towers, func, devices=None, use_vs=None)[source]

Run func on all GPUs (towers) and return the results.

Parameters:
  • towers (list[int]) – a list of GPU id.

  • func – a lambda to be called inside each tower

  • devices – a list of devices to be used. By default will use ‘/gpu:{tower}’

  • use_vs (list[bool]) – list of use_vs to passed to TowerContext

Returns:

List of outputs of func, evaluated on each tower.

class tensorpack.graph_builder.LeastLoadedDeviceSetter(worker_device, ps_devices)[source]

Bases: object

Helper class to assign variables on the least loaded ps-device.

Usage:

with tf.device(LeastLoadedDeviceSetter(...)):
    ...
__init__(worker_device, ps_devices)[source]
Parameters:
  • worker_device – the device to use for compute ops.

  • ps_devices – a list of device to use for Variable ops.

tensorpack.graph_builder.allreduce_grads(all_grads, average)[source]

All-reduce average the gradients among K devices. Results are broadcasted to all devices.

Parameters:
  • all_grads (K x N) – List of list of gradients. N is the number of variables.

  • average (bool) – average gradients or not.

Returns:

K x N – same as input, but each grad is replaced by the average over K devices.

tensorpack.graph_builder.aggregate_grads(all_grads, colocation=False, devices=None, average=True)[source]

Average the gradients.

Parameters:
  • all_grads (K x N x 2) – A list of K lists. Each of the list is a list of N (grad, var) tuples. The variables have to be the same across the K lists.

  • colocation (bool) – colocate gradient averaging on the device of the variable.

  • devices (list[str]) – assign the averaging to these device in round-robin. Cannot be used together with colocation.

  • average (bool) – do average or sum

Returns:

(N x 2) – A list of N (grad, var) tuples, where grad is averaged or summed over K.