mindspore.parallel.function.reshard

mindspore.parallel.function.reshard(tensor, layout)[源代码]

将张量从一种分布式排布转换成另一种分布式排布。其中，传入的layout需要为mindspore.parallel.Layout类型，可参考： mindspore.parallel.Layout 的描述。

说明

  • 在图模式下，可以利用此方法设置某个张量的并行切分策略，未设置的会自动通过策略传播方式配置。

  • 在PyNative模式下，可以利用此方法，对某个以图模式并行执行的Cell（即PyNative模式下使用了Cell.shard/F.shard的Cell）中的张量进行排布指定。

参数：

  • tensor (Tensor) - 待设置切分策略的张量。

  • layout (Layout) - 指定精准排布的方案，包括描述设备的排布（device_matrix）和设备矩阵的映射别名（alias_name）。

返回：

Tensor，与输入的tensor数学等价。

异常：

  • TypeError - 输入参数 tensor 不是mindspore.Tensor类型。

  • TypeError - 输入参数 layout 不是mindspore.parallel.Layout类型。

支持平台：

Ascend

样例：

说明

运行以下样例之前，需要配置好通信环境变量。

针对Ascend/GPU/CPU设备，推荐使用msrun启动方式，无第三方以及配置文件依赖。详见 msrun启动 。

该样例需要在8卡环境下运行。

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import ops, nn, Tensor, context, Layout
>>> from mindspore.parallel.function import reshard
>>> from mindspore.nn.utils import no_init_parameters
>>> from mindspore.parallel.auto_parallel import AutoParallel
>>> from mindspore.communication import init
>>> context.set_context(mode=ms.GRAPH_MODE)
>>> init()
>>> class Network(nn.Cell):
...     def __init__(self):
...         super().__init__()
...         self.matmul = ops.MatMul()
...         self.relu = ops.ReLU()
...     def construct(self, x, layout):
...         x = self.relu(x)
...         x_reshard = reshard(x, layout)
...         y = Tensor(np.ones(shape=(128, 128)), dtype=ms.float32)
...         x = self.matmul(x_reshard, y)
...         return x
>>> layout = Layout((4, 2), ("dp", "mp"))
>>> input_layout = layout("dp", "mp")
>>> with no_init_parameters():
...     net = Network()
>>> parallel_net = AutoParallel(net, parallel_mode='sharding_propagation')
>>> tensor = Tensor(np.ones(shape=(128, 128)), dtype=ms.float32)
>>> out = parallel_net(tensor, input_layout)