Python torch.allreduce方法代码示例

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce_type_error(self):
        """Test that the allreduce raises an error if different ranks try to
        send tensors of different type."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        # This test does not apply if there is only one worker.
        if size == 1:
            return

        # Same rank, different dimension
        dims = [17] * 3
        if rank % 2 == 0:
            tensor = torch.IntTensor(*dims)
        else:
            tensor = torch.FloatTensor(*dims)

        try:
            hvd.allreduce(tensor)
            assert False, 'hvd.allreduce did not throw error'
        except torch.FatalError:
            pass 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce_grad(self):
        """Test the correctness of the allreduce gradient."""
        hvd.init()
        size = hvd.size()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            tensor = torch.autograd.Variable(tensor, requires_grad=True)
            summed = hvd.allreduce(tensor, average=False)

            summed.backward(torch.ones([17] * dim))
            grad_out = tensor.grad.data.numpy()

            expected = np.ones([17] * dim) * size
            err = np.linalg.norm(expected - grad_out)
            self.assertLess(err, 0.00000001,
                            "gradient %s differs from expected %s, "
                            "error: %s" % (grad_out, expected, str(err))) 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def reduce_value(value, average, name):
    """
    Reduce the mean value of a tensor from all GPUs

    Parameters
    ----------
    value : torch.Tensor
        Value to be reduced
    average : bool
        Whether values will be averaged or not
    name : str
        Value name

    Returns
    -------
    value : torch.Tensor
        reduced value
    """
    return hvd.allreduce(value, average=average, name=name) 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def _distributed_update(self):
        if not is_distributed():
            return

        if is_horovod_available():
            import horovod.torch as hvd

            hvd.allreduce(self.track_num)
            hvd.allreduce(self.track_enc)
            hvd.allreduce(self.embed)
        else:
            distributed.all_reduce(self.track_num, op=distributed.ReduceOp.SUM)
            distributed.all_reduce(self.track_enc, op=distributed.ReduceOp.SUM)
            distributed.all_reduce(self.embed, op=distributed.ReduceOp.SUM)
            ws = distributed.get_world_size()
            self.track_num /= ws
            self.track_enc /= ws
            self.embed /= ws 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def metric_average(val, name):
    tensor = torch.FloatTensor([val])
    avg_tensor = hvd.allreduce(tensor, name=name)
    return avg_tensor.data[0] 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def update(self, val):
        self.sum += hvd.allreduce(val.cpu(), name=self.name)
        self.n += 1 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce(self):
        """Test that the allreduce correctly sums 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            summed = hvd.allreduce(tensor, average=False)
            multiplied = tensor * size
            max_difference = summed.data.sub(multiplied).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.IntTensor, torch.LongTensor,
                                      torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce_average(self):
        """Test that the allreduce correctly sums 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            averaged = hvd.allreduce(tensor, average=True)
            max_difference = averaged.data.sub(tensor).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.IntTensor, torch.LongTensor,
                                      torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce_inplace(self):
        """Test that the allreduce correctly sums 1D, 2D, 3D tensors."""
        hvd.init()
        size = hvd.size()
        dtypes = [torch.IntTensor, torch.LongTensor,
                  torch.FloatTensor, torch.DoubleTensor]
        if torch.cuda.is_available():
            dtypes += [torch.cuda.IntTensor, torch.cuda.LongTensor,
                       torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            tensor = tensor.type(dtype)
            multiplied = tensor * size
            hvd.allreduce_(tensor, average=False)
            max_difference = tensor.sub(multiplied).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.IntTensor, torch.LongTensor,
                                      torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce_multi_gpu(self):
        """Test that the allreduce works on multiple GPUs."""
        # Only do this test if there are GPUs available.
        if not torch.cuda.is_available():
            return

        hvd.init()
        local_rank = hvd.local_rank()
        size = hvd.size()

        iter = 0
        dtypes = [torch.cuda.IntTensor, torch.cuda.LongTensor,
                  torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            iter += 1
            torch.manual_seed(1234)
            tensor = torch.FloatTensor(*([17] * dim)).random_(-100, 100)
            device = local_rank * 2 + (iter + local_rank) % 2
            tensor = tensor.cuda(device).type(dtype)
            multiplied = tensor * size
            hvd.allreduce_(tensor, average=False)
            max_difference = tensor.sub(multiplied).max()

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if size <= 3 or dtype in [torch.cuda.IntTensor, torch.cuda.LongTensor]:
                threshold = 0
            elif size < 10:
                threshold = 1e-4
            elif size < 15:
                threshold = 5e-4
            else:
                break

            assert max_difference <= threshold, 'hvd.allreduce produces incorrect results' 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def test_horovod_allreduce_cpu_gpu_error(self):
        """Test that the allreduce raises an error if different ranks try to
        perform reduction on CPU and GPU."""
        # Only do this test if there are GPUs available.
        if not torch.cuda.is_available():
            return

        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        # This test does not apply if there is only one worker.
        if size == 1:
            return

        # Same rank, different dimension
        dims = [17] * 3
        if rank % 2 == 0:
            tensor = torch.cuda.FloatTensor(*dims)
        else:
            tensor = torch.FloatTensor(*dims)

        try:
            hvd.allreduce(tensor)
            assert False, 'hvd.allreduce did not throw error'
        except torch.FatalError:
            pass 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def metric_average(val, name):
    tensor = torch.tensor(val)
    avg_tensor = hvd.allreduce(tensor, name=name)
    return avg_tensor.item() 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def update(self, val):
        self.sum += hvd.allreduce(val.detach().cpu(), name=self.name)
        self.n += 1 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def _init_loaders(self):
        allreduce_batch_size = self.batch_size * self.batches_per_allreduce

        if hvd.local_rank() != 0:
            hvd.allreduce(torch.tensor(0), name='barrier')
        self.train_dataset = datasets.CIFAR10(
                root=self.dataset_path, download=(hvd.local_rank() == 0),
                train=True,
                transform=transforms.Compose([
                        transforms.RandomHorizontalFlip(),
                        transforms.RandomCrop(32, 4),
                        transforms.ToTensor(),
                        self.normalize,]))
        if hvd.local_rank() == 0:
            hvd.allreduce(torch.tensor(0), name='barrier')
        self.val_dataset = datasets.CIFAR10(
                root=self.dataset_path,
                train=False,
                transform=transforms.Compose([
                        transforms.ToTensor(),
                        self.normalize,]))
        self.train_loader = torch.utils.data.DataLoader(
                self.train_dataset, batch_size=allreduce_batch_size,
                shuffle=True, num_workers=8, pin_memory=True)
        self.val_loader = torch.utils.data.DataLoader(
                self.val_dataset, batch_size=allreduce_batch_size,
                shuffle=False, num_workers=8, pin_memory=True) 

# 需要导入模块: from horovod import torch [as 别名]
# 或者: from horovod.torch import allreduce [as 别名]
def reduce(self,
               tensor: torch.Tensor):

        if not torch.is_tensor(tensor):
            return tensor
        if is_distributed() and not self._no_reduce:
            if is_horovod_available():
                import horovod.torch as hvd

                # hvd's all_reduce applies average
                return hvd.allreduce(tensor)
            # pytorch's all_reduce does not applies average
            distributed.all_reduce(tensor, op=distributed.ReduceOp.SUM)
            return tensor / distributed.get_world_size()
        return tensor