mindspore.ops.ReduceScatter
- class mindspore.ops.ReduceScatter(op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP)[源代码]
规约并且分发指定通信组中的张量。
说明
在集合的所有过程中,Tensor必须具有相同的shape和格式。
- 参数:
op (str, 可选) - 指定用于元素的规约操作,如SUM和MAX。默认值:
ReduceOp.SUM
。group (str, 可选) - 要处理的通信组。默认值:
GlobalComm.WORLD_COMM_group
。
- 输入:
input_x (Tensor) - 输入Tensor,假设其形状为 \((N, *)\) ,其中 * 为任意数量的额外维度。N必须能够被rank_size整除,rank_size为当前通讯组里面的计算卡数量。
- 输出:
Tensor,数据类型与 input_x 一致,shape为 \((N/rank\_size, *)\) 。
- 异常:
TypeError - 如果 op 和 group 不是字符串。
ValueError - 如果输入的第一个维度不能被rank size整除。
- 支持平台:
Ascend
GPU
样例:
说明
运行以下样例之前,需要配置好通信环境变量。
针对Ascend设备,用户需要准备rank表,设置rank_id和device_id,详见 rank table启动 。
针对GPU设备,用户需要准备host文件和mpi,详见 mpirun启动 。
针对CPU设备,用户需要编写动态组网启动脚本,详见 动态组网启动 。
该样例需要在2卡环境下运行。
>>> import mindspore as ms >>> from mindspore import Tensor >>> from mindspore.communication import init >>> from mindspore.ops import ReduceOp >>> import mindspore.nn as nn >>> import mindspore.ops as ops >>> import numpy as np >>> >>> ms.set_context(mode=ms.GRAPH_MODE) >>> init() >>> class Net(nn.Cell): ... def __init__(self): ... super(Net, self).__init__() ... self.reducescatter = ops.ReduceScatter(ReduceOp.SUM) ... ... def construct(self, x): ... return self.reducescatter(x) ... >>> input_ = Tensor(np.ones([8, 8]).astype(np.float32)) >>> net = Net() >>> output = net(input_) >>> print(output) [[2. 2. 2. 2. 2. 2. 2. 2.] [2. 2. 2. 2. 2. 2. 2. 2.] [2. 2. 2. 2. 2. 2. 2. 2.] [2. 2. 2. 2. 2. 2. 2. 2.]]