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.

place_with_balance(op)[source]
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.

Parameters
  • 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”

Returns

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