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


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.

  • 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


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.


with tf.device(LeastLoadedDeviceSetter(...)):
__init__(worker_device, ps_devices)[source]
  • 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, mode='nccl')[source]

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

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

  • average (bool) – average gradients or not.

  • mode (str) – “nccl”, “collective”


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