Python Broadcast.apply方法代码示例

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def _data_parallel_master(self, intermediates):
        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""

        # Always using same "device order" makes the ReduceAdd operation faster.
        # Thanks to:: Tete Xiao (http://tetexiao.com/)
        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())

        to_reduce = [i[1][:2] for i in intermediates]
        to_reduce = [j for i in to_reduce for j in i]  # flatten
        target_gpus = [i[1].sum.get_device() for i in intermediates]

        sum_size = sum([i[1].sum_size for i in intermediates])
        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)

        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)

        outputs = []
        for i, rec in enumerate(intermediates):
            outputs.append((rec[0], _MasterMessage(*broadcasted[i * 2:i * 2 + 2])))

        return outputs 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def _data_parallel_master(self, intermediates):
        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())

        to_reduce = [i[1][:2] for i in intermediates]
        to_reduce = [j for i in to_reduce for j in i]  # flatten
        target_gpus = [i[1].sum.get_device() for i in intermediates]

        sum_size = sum([i[1].sum_size for i in intermediates])
        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)

        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)

        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)

        outputs = []
        for i, rec in enumerate(intermediates):
            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))

        return outputs 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def _data_parallel_master(self, intermediates):
        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""

        # Always using same "device order" makes the ReduceAdd operation faster.
        # Thanks to:: Tete Xiao (http://tetexiao.com/)
        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())

        to_reduce = [i[1][:2] for i in intermediates]
        to_reduce = [j for i in to_reduce for j in i]  # flatten
        target_gpus = [i[1].sum.get_device() for i in intermediates]

        sum_size = sum([i[1].sum_size for i in intermediates])
        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)

        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)

        outputs = []
        for i, rec in enumerate(intermediates):
            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))

        return outputs 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def batchnormtrain(input, mean, std, gamma, beta):
    r"""Applies Batch Normalization over a 3d input that is seen as a
    mini-batch.

    .. _encoding.batchnormtrain:

    .. math::

        y = \frac{x - \mu[x]}{ \sqrt{var[x] + \epsilon}} * \gamma + \beta

    Shape:
        - Input: :math:`(N, C)` or :math:`(N, C, L)`
        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)

    """
    return _batchnormtrain.apply(input, mean, std, gamma, beta) 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def forward(self, inputs, *targets, gathered=True, **kwargs):
        # input should be already scatterd
        # scattering the targets instead
        if gathered:
            if isinstance(inputs, (list, tuple)):
                inputs, _ = self.scatter(inputs, kwargs, self.device_ids)
            else:
                inputs, _ = self.scatter([inputs], kwargs, self.device_ids)
                # inputs = tuple(inputs_per_gpu[0] for inputs_per_gpu in inputs)

        if not self.device_ids:
            return self.module(inputs, *targets, **kwargs)

        targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            return self.module(inputs[0], *targets[0], **kwargs[0])

        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        # targets = tuple(targets_per_gpu[0] for targets_per_gpu in targets)
        outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs)
        return Reduce.apply(*outputs) / len(outputs) 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def data_parallel(f, input, params, mode, device_ids, output_device=None):
    assert isinstance(device_ids, list)
    if output_device is None:
        output_device = device_ids[0]

    if len(device_ids) == 1:
        return f(input, params, mode)

    params_all = Broadcast.apply(device_ids, *params.values())
    params_replicas = [{k: params_all[i + j*len(params)] for i, k in enumerate(params.keys())}
                       for j in range(len(device_ids))]

    replicas = [partial(f, params=p, mode=mode)
                for p in params_replicas]
    inputs = scatter([input], device_ids)
    outputs = parallel_apply(replicas, inputs)
    return gather(outputs, output_device) 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def _data_parallel_master(self, intermediates):
    """Reduce the sum and square-sum, compute the statistics, and broadcast it."""

    # Always using same "device order" makes the ReduceAdd operation faster.
    # Thanks to:: Tete Xiao (http://tetexiao.com/)
    intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())

    to_reduce = [i[1][:2] for i in intermediates]
    to_reduce = [j for i in to_reduce for j in i]  # flatten
    target_gpus = [i[1].sum.get_device() for i in intermediates]

    sum_size = sum([i[1].sum_size for i in intermediates])
    sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
    mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)

    broadcasted = Broadcast.apply(target_gpus, mean, inv_std)

    outputs = []
    for i, rec in enumerate(intermediates):
      outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))

    return outputs 

# 需要导入模块: from torch.nn.parallel._functions import Broadcast [as 别名]
# 或者: from torch.nn.parallel._functions.Broadcast import apply [as 别名]
def forward(self, inputs, *targets, **kwargs):
        if not self.device_ids:
            return self.module(inputs, *targets, **kwargs)

        is_target_scattered = kwargs.get('is_target_scattered', False)
        kwargs.pop('is_target_scattered', None)  # this key is unexpected

        if not is_target_scattered:
            targets, kwargs = self.scatter(targets, kwargs, self.device_ids)

        if len(self.device_ids) == 1:
            if is_target_scattered:
                targets = (targets,)
                kwargs = (kwargs,)
            return self.module(inputs, *targets[0], **kwargs[0])

        if is_target_scattered:
            targets = targets[0]

        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = _criterion_parallel_apply(replicas, inputs, targets)
        return Reduce.apply(*outputs) / len(outputs)