mmengine.dist.gather¶

mmengine.dist.gather(data, dst=0, group=None)[source]¶

Gather data from the whole group to dst process.

Note

Calling gather in non-distributed environment dose nothing and just returns a list containing data itself.

Note

NCCL backend does not support gather.

Note

Unlike PyTorch torch.distributed.gather, gather() in MMEngine does not pass in an empty list gather_list and returns the gather_list directly, which is more convenient. The difference between their interfaces is as below:

MMEngine: gather(data, dst, group) -> gather_list
PyTorch: gather(data, gather_list, dst, group) -> None

Parameters:

data (Tensor) – Tensor to be gathered. CUDA tensor is not supported.
dst (int) – Destination rank. Defaults to 0.
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used. Defaults to None.

Returns:

dst process will get a list of tensor gathering from the whole group. Other process will get a empty list. If in non-distributed environment, just return a list containing data itself.

Return type:

list[Tensor]

Examples

>>> import torch
>>> import mmengine.dist as dist

>>> # non-distributed environment
>>> data = torch.arange(2, dtype=torch.int64)
>>> data
tensor([0, 1])
>>> output = dist.gather(data)
>>> output
[tensor([0, 1])]

>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> data = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> data
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> output = dist.gather(data)
>>> output
[tensor([1, 2]), tensor([3, 4])]  # Rank 0
[]  # Rank 1