【onnx】——ScatterND操作

最新推荐文章于 2024-01-31 21:08:23 发布

农夫山泉2号

最新推荐文章于 2024-01-31 21:08:23 发布

阅读量8.9k

点赞数 6

分类专栏：嵌入式AI 文章标签： tensorrt onnx

本文链接：https://blog.csdn.net/u011622208/article/details/115013246

版权

嵌入式AI 专栏收录该内容

158 篇文章 52 订阅

订阅专栏

scatterND，ONNX，pytorch

pytorch中的操作转onnx时的一个操作。其实就是一个索引的意思。比如下面的一段代码：

x = torch.randn(20, 200, 200)
y = torch.randn(10, 200, 200)
x[0:10, :, :] += y

这样一段代码转onnx的时候就会有scatterND算子。

scatterND算子的定义：https://github.com/onnx/onnx/blob/master/docs/Changelog.md#ScatterND-11

在这里插入图片描述

其输入有3个，输出只有1个。

output的计算规则

  output = np.copy(data)
  update_indices = indices.shape[:-1]
  for idx in np.ndindex(update_indices):
      output[indices[idx]] = updates[idx]

可以通过以下的2个例子来理解算子的计算
Example 1:

  data    = [1, 2, 3, 4, 5, 6, 7, 8]
  indices = [[4], [3], [1], [7]]
  updates = [9, 10, 11, 12]
  output  = [1, 11, 3, 10, 9, 6, 7, 12]

Example 2:

  data    = [[[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
             [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
             [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]],
             [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]]
  indices = [[0], [2]]
  updates = [[[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
             [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]]]
  output  = [[[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
             [[1, 2, 3, 4], [5, 6, 7, 8], [8, 7, 6, 5], [4, 3, 2, 1]],
             [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4]],
             [[8, 7, 6, 5], [4, 3, 2, 1], [1, 2, 3, 4], [5, 6, 7, 8]]]

可以看到在一维的时候是按indices替换data。多维的时候，是output的第一维替换output[indices[idx]]