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

class
tensorpack.graph_builder.
LeastLoadedDeviceSetter
(worker_device, ps_devices)[source]¶ Bases:
object
Helper class to assign variables on the least loaded psdevice.
Usage:
with tf.device(LeastLoadedDeviceSetter(...)): ...

tensorpack.graph_builder.
allreduce_grads
(all_grads, average)[source]¶ Allreduce 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 roundrobin. 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.