mindspore.ops.tensordump

查看源文件
mindspore.ops.tensordump(file_name, tensor, mode='out')[源代码]

将Tensor保存为Numpy的npy格式的文件。

在并行的场景下,该算子会保存不同计算卡上的数据分片。

在昇腾平台的静态图并行模式下,代码 算子A –> 算子B 可能会被编译为 算子A –> 重排算子 –> 算子B。

注:重排算子指在静态图并行场景中,由于设备间通信,算子切分策略而引入的算子。

对于 算子A –> 算子B 的情况,算子A计算结果的输出等于算子B的输入。

然而对于算子A –> 重排算子 –> 算子B这样的情况,由于重排算子的插入,导致算子A的输出不再等价于算子B的输入。

假设一种情况是算子A的计算结果既作为算子B的输入,也作为tensordump算子的输入。则在该情况下, 通过设置mode参数可以实现不同的保存数据的需求:

  • 如果mode参数设置为'out',保存的数据仅包含算子A的输出分片。

  • 如果mode参数设置为'all',保存的数据将包含算子A的输出分片以及算子B的输入分片。

  • 如果mode参数设置为'in',保存的数据将仅包含算子B的输入分片。

当mode参数被配置为'all'或'in'时,生成的输入分片所对应的npy文件命名格式为:id_fileName_cNodeID_dumpMode_rankID.npy。

当mode参数被配置为'all'或'out'时,生成的输出分片所对应的npy文件命名格式为:id_filename.npy。

  • id:一个自增的ID。

  • fileName:参数file_name的值 (若该参数传入时是一个使用者指定的路径,则fileName的值为路径的最后一级)。

  • cNodeID:该tensordump节点在step_parallel_end.ir文件中的节点编号

  • dumpMode:mode参数的值。

  • rankID:逻辑卡号。

说明

  • 当前该算子不支持在控制流中使用。

  • 如果当前的并行模式为STAND_ALONE,参数mode只能设置为'out'。

  • 如使用该算子时不设置参数mode,其默认值为'out'。

参数:
  • file_name (str) - npy文件的保存路径。

  • tensor (Tensor) - 输入的张量。

  • mode (str,可选) - 控制tensordump行为模式的参数,可选的值为 ['out', 'in', 'all'] 中的一个,默认值: out

异常:
  • TypeError - file_name 不是一个str类型。

  • TypeError - tensor 不是一个Tensor类型。

  • TypeError - mode 不是一个str类型。

  • ValueError - mode 的值不是 ['out', 'in', 'all'] 之中的一个。

支持平台:

Ascend

样例:

说明

使用msrun命令运行下面的例子:msrun –worker_num=2 –local_worker_num=2 –master_port=11450 –log_dir=msrun_log –join=True –cluster_time_out=300 tensordump_example.py

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import nn, Tensor, ops, context
>>> from mindspore.ops import operations as P
>>> from mindspore.communication import init, get_rank
>>> init()
>>> rank_id = get_rank()
>>> dump_path = f'dumps/rank_{rank_id}/mul1_mul2.npy'
>>> class Net(nn.Cell):
...     def __init__(self, strategy1, strategy2):
...         super(Net, self).__init__()
...         self.matmul1 = P.MatMul().shard(strategy1)
...         self.matmul2 = P.MatMul().shard(strategy2)
...
...     def construct(self, x, y, b):
...         out1 = self.matmul1(x, y)
...         ops.tensordump(dump_path, out1, 'all')
...         out2 = self.matmul2(out1, b)
...         return out2
...
>>> ms.set_context(mode=ms.GRAPH_MODE, save_graphs=2)
>>> context.set_auto_parallel_context(parallel_mode='semi_auto_parallel', full_batch=True)
>>> strategy1 = ((1, 2), (2, 1))
>>> strategy2 = ((1, 2), (2, 1))
>>> net = Net(strategy1, strategy2)
>>> x = Tensor(0.1 * np.random.randn(64, 64).astype(np.float32))
>>> y = Tensor(0.1 * np.random.randn(64, 64).astype(np.float32))
>>> b = Tensor(0.1 * np.random.randn(64, 64).astype(np.float32))
>>> out = net(x, y, b)
>>> print(f"out shape is: {out.shape}")
>>> matmul1_output_slice = np.load('0_mul1_mul2.npy')                       # load matmul1's output slice
>>> matmul2_input_slice = np.load('1_mul1_mul2_CNode_64_all_rank_0.npy')    # load matmul2's input slice