Python PyTorch all_reduce用法及代碼示例

本文簡要介紹python語言中 torch.distributed.all_reduce 的用法。

用法:

torch.distributed.all_reduce(tensor, op=<ReduceOp.SUM: 0>, group=None, async_op=False)

參數：

• tensor(Tensor) -集體的輸入和輸出。該函數就地運行。

• op(可選的) -torch.distributed.ReduceOp 枚舉中的值之一。指定用於按元素減少的操作。

• group(ProcessGroup,可選的) -要處理的流程組。如果沒有，將使用默認進程組。

• async_op(bool,可選的) -此操作是否應該是異步操作

返回：

異步工作句柄，如果 async_op 設置為 True。無，如果不是 async_op 或者如果不是該組的一部分

以這樣的方式減少所有機器上的張量數據，從而得到最終結果。

在調用 tensor 之後，所有進程將按位相同。

支持複雜的張量。

例子

>>> # All tensors below are of torch.int64 type.
>>> # We have 2 process groups, 2 ranks.
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> tensor
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_reduce(tensor, op=ReduceOp.SUM)
>>> tensor
tensor([4, 6]) # Rank 0
tensor([4, 6]) # Rank 1

>>> # All tensors below are of torch.cfloat type.
>>> # We have 2 process groups, 2 ranks.
>>> 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_reduce(tensor, op=ReduceOp.SUM)
>>> tensor
tensor([4.+4.j, 6.+6.j]) # Rank 0
tensor([4.+4.j, 6.+6.j]) # Rank 1