mmengine.dist.collect_results_gpu

mmengine.dist.collect_results_gpu(result_part, size)

Collect results under gpu mode.

On gpu mode, this function will encode results to gpu tensors and use gpu communication for results collection.

Parameters:

• result_part (list[object]) – 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.

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_gpu(data, size)
>>> output
['foo', 24, {1: 2}, {'a': 'b'}]  # rank 0
None  # rank 1