mindspore.ops.NeighborExchangeV2
- class mindspore.ops.NeighborExchangeV2(send_rank_ids, recv_rank_ids, send_lens, recv_lens, data_format, group=GlobalComm.WORLD_COMM_GROUP)[源代码]
NeighborExchangeV2是一个集合通讯操作。
将数据从本地rank发送到 send_rank_ids 中指定的rank,同时从 recv_rank_ids 接收数据。请参考下方教程样例了解具体的数据是如何在相邻设备间交换的。
说明
要求全连接配网,每台设备具有相同的vlan id,ip和mask在同一子网,请查看 分布式集合通信原语注意事项 。
- 参数:
send_rank_ids (list(int)) - 指定发送数据的rank。8个rank_id分别代表8个方向上的数据要向哪个rank发送,如果某个方向上不发送数据,则设为-1。
recv_rank_ids (list(int)) - 指定接收数据的rank。8个rank_id分别代表8个方向上的数据要从哪个rank接收,如果某个方向上不接收数据,则设为-1。
send_lens (list(int)) - 指定 send_rank_ids 发送数据的长度,4个数字分别代表[send_top, send_bottom, send_left, send_right]4个方向上的长度。
recv_lens (list(int)) - 指定 recv_rank_ids 接收数据的长度,4个数字分别代表[recv_top, recv_bottom, recv_left, recv_right]4个方向上的长度。
data_format (str) - 数据格式,现在只支持NCHW。
group (str, 可选) - 工作的通信组。默认值:
GlobalComm.WORLD_COMM_GROUP
(即Ascend平台为"hccl_world_group"
,GPU平台为"nccl_world_group"
)。
- 输入:
input_x (Tensor) - 交换前的输入Tensor,其shape为 \((N, C, H, W)\)。
- 输出:
数据交换后的输出Tensor,如果输入的shape是 \((N, C, H, W)\) ,则输出shape为 \((N, C, H+recv\_top+recv\_bottom, W+recv\_left+recv\_right)\) 。
- 异常:
TypeError - 如果 group 不是一个string或者 send_rank_ids 、 recv_rank_ids 、 send_lens 和 recv_lens 中任意一个不是一个list。
ValueError - 如果 send_rank_ids 或者 recv_rank_ids 存在小于-1的值或者存在重复值。
ValueError - 如果 send_lens 或者 recv_lens 存在小于零的值。
ValueError - 如果 data_format 不是"NCHW"。
- 支持平台:
Ascend
样例:
>>> import os >>> import mindspore as ms >>> from mindspore.communication import init >>> import mindspore.nn as nn >>> from mindspore import ops >>> import numpy as np >>> >>> class Net0(nn.Cell): ... def __init__(self): ... super(Net0, self).__init__() ... self.neighbor_exchangev2 = ops.NeighborExchangeV2(send_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1], ... send_lens=[0, 1, 0, 0], ... recv_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1], ... recv_lens=[0, 1, 0, 0], data_format="NCHW") ... ... def construct(self, x): ... out = self.neighbor_exchangev2(x) ... return out ... class Net1(nn.Cell): ... def __init__(self): ... super(Net1, self).__init__() ... self.neighbor_exchangev2 = ops.NeighborExchangeV2(send_rank_ids=[0, -1, -1, -1, -1, -1, -1, -1], ... send_lens=[1, 0, 0, 0], ... recv_rank_ids=[0, -1, -1, -1, -1, -1, -1, -1], ... recv_lens=[1, 0, 0, 0], data_format="NCHW") ... ... def construct(self, x): ... out = self.neighbor_exchangev2(x) ... return out >>> >>> ms.set_context(mode=ms.GRAPH_MODE) >>> init() >>> rank_id = int(os.getenv("RANK_ID")) >>> if (rank_id % 2 == 0): >>> input_x = ms.Tensor(np.ones([1, 1, 2, 2]), dtype = ms.float32) >>> net = Net0() >>> output = net(input_x) >>> print(output) >>> else: >>> input_x = ms.Tensor(np.ones([1, 1, 2, 2]) * 2, dtype = ms.float32) >>> net = Net1() >>> output = net(input_x) >>> print(output) [[[[1. 1.], [1. 1.], [2. 2.]]]]