mmengine.dist.gather_object¶

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

Gathers picklable objects from the whole group in a single process. Similar to gather(), but Python objects can be passed in. Note that the object must be picklable in order to be gathered.

Note

NCCL backend does not support gather_object.

Note

Unlike PyTorch torch.distributed.gather_object, gather_object() 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_object(data, dst, group) -> gather_list
PyTorch: gather_object(data, gather_list, data, group) -> None

Parameters:

data (Any) – Input object. Must be picklable.
dst (int) – Destination rank. Defaults to 0.
group (ProcessGroup | None) – (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Defaults to None.

Returns:

list[Any]. On the dst rank, return gather_list which contains the output of the collective.

Return type:

List[Any] | None

Examples

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

>>> # non-distributed environment
>>> data = ['foo', 12, {1: 2}]  # any picklable object
>>> gather_objects = dist.gather_object(data[dist.get_rank()])
>>> output
['foo']

>>> # distributed environment
>>> # We have 3 process groups, 3 ranks.
>>> dist.gather_object(gather_objects[dist.get_rank()], dst=0)
>>> output
['foo', 12, {1: 2}]  # Rank 0
None  # Rank 1
None  # Rank 2