本文介绍了MindSpore中的tensor_scatter_add算子的用法,可以对一个多维的tensor在指定的index上面进行加和操作。在PyTorch中虽然也有一个叫scatter_add的算子,但是本质上来说两者是完全不一样的操作。
技术背景
在MindSpore的ops下实现了一个tensor_scatter_add算子。这个算子的作用为,例如给定一个shape为(1,2,3,4)的原始tensor,因为这个tensor有4个维度,所以我们每次去提取一个tensor元素的时候,就需要4个索引。那么假如说我们要提取这个tensor中的5个元素,那么需要用到的索引tensor的shape应该为(5,4)。这样一来就可以提取得到5个元素,然后做一个add的操作,给定一个目标的tensor,它的shape为(5,),我们就可以把这个目标tensor按照索引tensor,加到原始的tensor里面去,这就是tensor_scatter_add算子的作用。但是在PyTorch中没有这个算子的实现,只有scatter_add和index_add,但是这三个算子的作用是完全不一样的,接下来用代码示例演示一下。
代码实现
首先看一个mindspore的tensor_scatter_add算子,其官方文档的介绍是这样的:
以下是一个代码实现的示例:
在MindSpore中该算子支持的是挺好的,因为在索引里面其实存在重复索引,并行计算的话可能存在race condition
的问题,而MindSpore中还是可以做到正确的加和。在PyTorch中有一个scatter_add算子,但是跟MindSpore里面的tensor_scatter_add完全是两个不一样的算子,以下是一个示例:
他这个scatter_add算子就不是为了这种场景而设计的。
还有另外一个index_add,但是用法也不太相似:
不过如果想要用PyTorch实现这个功能的话,也不是没有办法。还是以这个例子来说,因为原始tensor有4个维度,所以索引到每一个元素需要4个索引编号。在PyTorch中可以直接这样操作(不建议!不建议!不建议!
):
这个例子就可以看出来,直接使用索引进行加和的话,结果是不对的。存在相同索引的情况下,不同的操作之间有可能相互覆盖。所以,非常的不建议这么操作,除非能够确保索引tensor都是唯一的。
总结概要
本文介绍了MindSpore中的tensor_scatter_add算子的用法,可以对一个多维的tensor在指定的index上面进行加和操作。在PyTorch中虽然也有一个叫scatter_add的算子,但是本质上来说两者是完全不一样的操作。
作者ID:DechinPhy