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)

Bases: tensorpack.graph_builder.training.GraphBuilder

__init__(towers)

Parameters

towers (list[int]) – list of GPU ids.

static call_for_each_tower(towers, func, devices=None, use_vs=None)

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)

Bases: object

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

Usage:

with tf.device(LeastLoadedDeviceSetter(...)):
    ...

__init__(worker_device, ps_devices)

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)

tensorpack.graph_builder.allreduce_grads(all_grads, average, mode='nccl')

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.