Python PyTorch all_gather用法及代码示例

本文简要介绍python语言中 torch.distributed.all_gather 的用法。

用法:

torch.distributed.all_gather(tensor_list, tensor, group=None, async_op=False)

参数：

• tensor_list(list[Tensor]) -输出列表。它应该包含 correctly-sized 张量，用于集体的输出。

• tensor(Tensor) -从当前进程广播的张量。

• group(ProcessGroup,可选的) -要处理的流程组。如果没有，将使用默认进程组。

• async_op(bool,可选的) -此操作是否应该是异步操作

返回：

异步工作句柄，如果 async_op 设置为 True。无，如果不是 async_op 或者如果不是该组的一部分

将整个组的张量收集到一个列表中。

支持复杂的张量。

例子

>>> # All tensors below are of torch.int64 dtype.
>>> # We have 2 process groups, 2 ranks.
>>> tensor_list = [torch.zeros(2, dtype=torch.int64) for _ in range(2)]
>>> tensor_list
[tensor([0, 0]), tensor([0, 0])] # Rank 0 and 1
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> tensor
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_gather(tensor_list, tensor)
>>> tensor_list
[tensor([1, 2]), tensor([3, 4])] # Rank 0
[tensor([1, 2]), tensor([3, 4])] # Rank 1

>>> # All tensors below are of torch.cfloat dtype.
>>> # We have 2 process groups, 2 ranks.
>>> tensor_list = [torch.zeros(2, dtype=torch.cfloat) for _ in range(2)]
>>> tensor_list
[tensor([0.+0.j, 0.+0.j]), tensor([0.+0.j, 0.+0.j])] # Rank 0 and 1
>>> tensor = torch.tensor([1+1j, 2+2j], dtype=torch.cfloat) + 2 * rank * (1+1j)
>>> tensor
tensor([1.+1.j, 2.+2.j]) # Rank 0
tensor([3.+3.j, 4.+4.j]) # Rank 1
>>> dist.all_gather(tensor_list, tensor)
>>> tensor_list
[tensor([1.+1.j, 2.+2.j]), tensor([3.+3.j, 4.+4.j])] # Rank 0
[tensor([1.+1.j, 2.+2.j]), tensor([3.+3.j, 4.+4.j])] # Rank 1