Shortcuts

mmengine.dist.all_gather_object

mmengine.dist.all_gather_object(data, group=None)[source]

Gather picklable objects from the whole group into a list. Similar to all_gather(), but Python objects can be passed in. Note that the object must be picklable in order to be gathered.

Note

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

Note

Unlike PyTorch torch.distributed.all_gather_object, all_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: all_gather_object(data, group) -> gather_list

  • PyTorch: all_gather_object(gather_list, data, group) -> None

Parameters
  • data (Any) – Pickable Python object to be broadcast from current process.

  • group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used. Defaults to None.

Returns

Return a list containing data from the whole group if in distributed environment, otherwise a list only containing data itself.

Return type

list[Tensor]

Note

For NCCL-based process groups, internal tensor representations of objects must be moved to the GPU device before communication starts. In this case, the used device is given by torch.cuda.current_device() and it is the user’s responsibility to ensure that this is correctly set so that each rank has an individual GPU, via torch.cuda.set_device().

Examples

>>> import torch
>>> import mmengine.dist as dist
>>> # non-distributed environment
>>> data = ['foo', 12, {1: 2}]  # any picklable object
>>> gather_objects = dist.all_gather_object(data[dist.get_rank()])
>>> output
['foo']
>>> # distributed environment
>>> # We have 3 process groups, 3 ranks.
>>> output = dist.all_gather_object(data[dist.get_rank()])
>>> output
['foo', 12, {1: 2}]  # Rank 0
['foo', 12, {1: 2}]  # Rank 1
['foo', 12, {1: 2}]  # Rank 2
Read the Docs v: v0.8.3
Versions
latest
stable
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.