Class SyncMultiGPUReplicatedBuilder

tensorpack/graph_builder/training.py:196–353

Data-parallel training in "replicated" mode, where each GPU contains a replicate of the whole model. It will build one tower on each GPU under its own variable scope. Each gradient update is averaged or summed across or GPUs through NCCL. It is an equivalent of ``--variable_upd

class SyncMultiGPUReplicatedBuilder(DataParallelBuilder):
    """
    Data-parallel training in "replicated" mode,
    where each GPU contains a replicate of the whole model.
    It will build one tower on each GPU under its own variable scope.
    Each gradient update is averaged or summed across or GPUs through NCCL.

    It is an equivalent of ``--variable_update=replicated`` in
    `tensorflow/benchmarks <https://github.com/tensorflow/benchmarks>`_.
    """

    def __init__(self, towers, average, mode):
        super(SyncMultiGPUReplicatedBuilder, self).__init__(towers)
        self._average = average
        assert mode in ['nccl', 'cpu', 'hierarchical', 'gpu', 'collective'], mode
        self._mode = mode

        if self._mode == 'hierarchical' and len(towers) != 8:
            raise ValueError("mode='hierarchical' require 8 GPUs.")

    def call_for_each_tower(self, tower_fn):
        """
        Call the function `tower_fn` under :class:`TowerContext` for each tower.

        Returns:
            a list, contains the return values of `tower_fn` on each tower.
        """
        # if tower_fn returns [(grad, var), ...], this returns #GPU x #VAR x 2
        return DataParallelBuilder.build_on_towers(
            self.towers,
            tower_fn,
            # use no variable scope for the first tower
            use_vs=[False] + [True] * (len(self.towers) - 1))

    def build(self, grad_list, get_opt_fn):
        """
        Reduce the gradients, apply them with the optimizer,
        and set self.grads to #GPU number of lists of (g, v), containing the all-reduced gradients on each device.

        Args:
            grad_list ([[(grad, var), ...], ...]): #GPU lists to be reduced. Each is the gradients computed on each GPU.
            get_opt_fn (-> tf.train.Optimizer): callable which returns an optimizer

        Returns:
            (tf.Operation, tf.Operation)

            1. the training op.

            2. the op which sync variables from GPU 0 to other GPUs.
                It has to be run before the training has started.
                And you can optionally run it later to sync non-trainable variables.
        """
        assert len(grad_list) == len(self.towers)
        raw_devices = ['/gpu:{}'.format(k) for k in self.towers]

        DataParallelBuilder._check_grad_list(grad_list)

        dtypes = {x[0].dtype.base_dtype for x in grad_list[0]}