知乎介绍参考添加链接描述
完整测试框架参考本人仓库
添加链接描述
gather算子的onnx定义参考添加链接描述,该算子的主要变换参考下图:
这里我们不妨以input = [A, dimsize, D], indices = [B,C], axis = 1举例子,此时对应的output形状是[A,B,C,D],并且根据gather算子定义,我们知道output[i,j,k,s] = input[i,indices[j,k], s],下面我们引入一些定义:
outputSize = ABCD, inputSize = AD×dimsize, indSize = BC。
那么此时我们可以看出,首先input和output共享一个othersize = AD,我们可以这么思考,对于output来说,可以把output分成两部分,看成一个[othersize, indSize]的矩阵,同样的,input也可以分成两部分,看成一个[othersize, dimsize]的矩阵,因此output和input共享了othersize部分的数据,后面的主要数学变换集中在后面部分。
除此之外,不管axis在哪,对于input和output来说,对应的stride是一样的,这里的stride表示axis对应的步