megengine.functional.distributed.all_reduce_min

all_reduce_min(inp, group=WORLD, device=None)

在指定组中以 min 操作来对张量进行规约操作。

注解

inp 张量在整个组的所有进程中必须形状相同。

参数

inp (Tensor) – 需要规约的张量

关键字参数

• group (Group or sequence of ints) – 需要处理的组。默认值：“WORLD”. WORLD 组包含所有可用的进程。可以以进程序号构建一个列表以创建一个新的组。

• device (Tensor.device) – 执行此操作的设备。默认值： None. None 表示以 inp 的设备作为执行设备。特别地， GPU 设备可以通过在设备名后的冒号后面添加一个数字来分配要执行的 cuda 流，而 :0 表示 GPU 的默认流，不指定流则表示使用默认流。

返回类型

Tensor

返回

对组中的的值进行 min 操作所得到的张量。输出张量的形状必须与 inp 相同，并且输出张量在整个组的所有进程中都是完全相同的。

实际案例

>>> # We execute all_reduce_min on rank 0 and rank 1
>>> input = F.arange(2) + 1 + 2 * rank 
>>> input  
Tensor([1. 2.], device=xpux:0) # Rank 0
Tensor([3. 4.], device=xpux:0) # Rank 1
>>> F.distributed.all_reduce_min(input, group=[0, 1]) 
Tensor([1. 2.], device=xpux:0) # Rank 0
Tensor([1. 2.], device=xpux:0) # Rank 1

>>> # We execute all_reduce_min with on gpu0 with cuda stream 1
>>> megengine.set_default_device("gpu0") 
>>> input = F.arange(2) + 1 + 2 * rank 
>>> input  
Tensor([1. 2.], device=gpu0:0) # Rank 0
Tensor([3. 4.], device=gpu0:0) # Rank 1
>>> F.distributed.all_reduce_min(input, device="gpu0:1") 
Tensor([1. 2.], device=xpux:0) # Rank 0
Tensor([1. 2.], device=xpux:0) # Rank 1

上一页

megengine.functional.distributed.all_reduce_max

下一页

megengine.functional.distributed.all_reduce_sum