Ejemplo n.º 1
0
def broadcast_parameters(params, root_rank=0):
    """
    Broadcasts the parameters from root rank to all other processes.
    Typical usage is to broadcast the `Module.get_params()` or the
    `Block.collect_params()`.
    Arguments:
        params: One of the following:
            - dict of parameters to broadcast
            - ParameterDict to broadcast
        root_rank: The rank of the process from which parameters will be
                   broadcasted to all other processes.
    """
    tensors = []
    if isinstance(params, dict):
        tensors = [p for _, p in sorted(params.items())]
    elif isinstance(params, mx.gluon.parameter.ParameterDict):
        for _, p in sorted(params.items()):
            try:
                tensors.append(p.data())
            except mx.gluon.parameter.DeferredInitializationError:
                # Inject wrapper method with post-initialization broadcast to
                # handle parameters with deferred initialization
                global parameter_index
                byteps_declare_tensor(p.data(),
                                      "parameter_" + str(parameter_index))
                new_init = _append_broadcast_init(p, root_rank,
                                                  parameter_index)
                parameter_index += 1
                p._init_impl = types.MethodType(new_init, p)
    else:
        raise ValueError('invalid params of type: %s' % type(params))

    # Run tensor initilization
    for i in range(len(tensors)):
        byteps_declare_tensor(tensors[i], "parameter_" + str(parameter_index))
        # Broadcast is implemented as push + pull in BytePS
        # To broadcast: we should zero-out all non-root tensors, and disable push_pull average
        if rank() != root_rank:
            tensors[i].__imul__(0)
        byteps_push_pull(tensors[i],
                         version=0,
                         priority=0,
                         name="parameter_" + str(parameter_index),
                         is_average=False)
        parameter_index += 1

    # Make sure tensors pushed to MXNet engine get processed such that all
    # workers are synced before starting training.
    for tensor in tensors:
        tensor.wait_to_read()
Ejemplo n.º 2
0
    def _init_params(self):
        tensors = []
        for param in self._params_to_init:
            if param._deferred_init:
                tensors.append(param)
            else:
                param_arrays = param._check_and_get(param._data, list)
                idx = self._param2idx[param.name]

                if rank() != self.root_rank:
                    param_arrays[0].__imul__(0)
                byteps_push_pull(param_arrays[0],
                                 version=0,
                                 priority=0,
                                 name="parameter_" + str(idx),
                                 is_average=False)

        self._params_to_init = tensors
Ejemplo n.º 3
0
def broadcast_parameters(params, root_rank=0):
    """
    Broadcasts the parameters from root rank to all other processes.
    Typical usage is to broadcast the `Module.get_params()`.

    Arguments:
        params: dict of parameters to broadcast
        root_rank: The rank of the process from which parameters will be
                   broadcasted to all other processes.
    """
    global parameter_index

    if isinstance(params, dict):
        tensors = [p for _, p in sorted(params.items())]

        # Run tensor initilization
        for i in range(len(tensors)):
            byteps_declare_tensor(tensors[i],
                                  "parameter_" + str(parameter_index))
            # Broadcast is implemented as push + pull in BytePS
            # To broadcast: we should zero-out all non-root tensors, and disable push_pull average
            if rank() != root_rank:
                tensors[i].__imul__(0)
            byteps_push_pull(tensors[i],
                             version=0,
                             priority=0,
                             name="parameter_" + str(parameter_index),
                             is_average=False)
            parameter_index += 1

        # Make sure tensors pushed to MXNet engine get processed such that all
        # workers are synced before starting training.
        for tensor in tensors:
            tensor.wait_to_read()

    elif isinstance(params, mx.gluon.parameter.ParameterDict):
        raise TypeError("For gluon users, you should not call this function. "
                        "DistributedTrainer will broadcast all parameters at "
                        "the first training step.")

    else:
        raise ValueError('Invalid params of type: %s' % type(params))