mmengine.dist.collect_results_cpu¶

mmengine.dist.collect_results_cpu(result_part, size, tmpdir=None)[source]¶

Collect results under cpu mode.

On cpu mode, this function will save the results on different gpus to tmpdir and collect them by the rank 0 worker.

Parameters

result_part (list) – Result list containing result parts to be collected. Each item of result_part should be a picklable object.
size (int) – Size of the results, commonly equal to length of the results.
tmpdir (str | None) – Temporal directory for collected results to store. If set to None, it will create a random temporal directory for it. Defaults to None.

Returns

The collected results.

Return type

list or None

Examples

>>> # distributed environment
>>> # We have 2 process groups, 2 ranks.
>>> import mmengine.dist as dist
>>> if dist.get_rank() == 0:
        data = ['foo', {1: 2}]
    else:
        data = [24, {'a': 'b'}]
>>> size = 4
>>> output = dist.collect_results_cpu(data, size)
>>> output
['foo', 24, {1: 2}, {'a': 'b'}]  # rank 0
None  # rank 1