mmengine.dist.gather — mmengine 0.10.7 documentation (original) (raw)

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 containingdata 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