Class DistributedReplicatedBuilder

tensorpack/graph_builder/distributed.py:135–374

Distributed replicated training. Each worker process builds the same model on one or more GPUs. Gradients across GPUs are averaged within the worker, and get synchronously applied to the global copy of variables located on PS. Then each worker copy the latest variables from PS b

class DistributedReplicatedBuilder(DataParallelBuilder, DistributedBuilderBase):
    """
    Distributed replicated training.
    Each worker process builds the same model on one or more GPUs.
    Gradients across GPUs are averaged within the worker,
    and get synchronously applied to the global copy of variables located on PS.
    Then each worker copy the latest variables from PS back to local.

    It is an equivalent of ``--variable_update=distributed_replicated`` in
    `tensorflow/benchmarks <https://github.com/tensorflow/benchmarks>`_.
    Note that the performance of this trainer is still not satisfactory,
    and TensorFlow team is not actively maintaining distributed training features.
    Check :class:`HorovodTrainer` and
    `ResNet-Horovod <https://github.com/tensorpack/benchmarks/tree/master/ResNet-Horovod>`_
    for better distributed training support.

    Note:
        1. Gradients are not averaged across workers, but applied to PS variables
           directly (either with or without locking depending on the optimizer).
        2. Some details about collections: all variables created inside tower
           will become local variables,
           and a clone will be made in global variables for all trainable/model variables.

    Example:

        .. code-block:: python

            # Create the server object like this:
            hosts = ['host1.com', 'host2.com']
            cluster_spec = tf.train.ClusterSpec({
                'ps': [h + ':2222' for h in hosts],
                'worker': [h + ':2223' for h in hosts]
            })
            server = tf.train.Server(
                cluster_spec, job_name=args.job, task_index=args.task,
                config=get_default_sess_config())
            # initialize trainer with this server object

        .. code-block:: none

            # Start training like this:
            (host1)$ ./train.py --job worker --task 0
            (host1)$ CUDA_VISIBLE_DEVICES= ./train.py --job ps --task 0
            (host2)$ ./train.py --job worker --task 1
            (host2)$ CUDA_VISIBLE_DEVICES= ./train.py --job ps --task 1
    """

    def __init__(self, towers, server):
        """
        Args:
            towers (list[int]): list of GPU ids.
            server (tf.train.Server): the server with ps and workers.
                job_name must be 'worker'.
        """
        DataParallelBuilder.__init__(self, towers)
        DistributedBuilderBase.__init__(self, server)

        self.is_chief = (self.task_index == 0)