飞桨Paddle API index_add 详解

沿着指定轴 axis 将 index 中指定位置的 x 与 value 相加，并写入到结果 Tensor 中的对应位置。这里 index 是一个 1-D Tensor。除 axis 轴外，返回的 Tensor 其余维度大小和输入 x 相等， axis 维度的大小等于 index 的大小。

官方文档：index_add-API文档-PaddlePaddle深度学习平台

我们还是通过一个代码示例来学习：

x = paddle.ones([5, 3])
value = paddle.to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=paddle.float32)
index = paddle.to_tensor([0, 4, 2])
print(x)x = paddle.index_add(x, index, 0, value)
print(x)

 输出

Tensor(shape=[5, 3], dtype=float32, place=Place(cpu), stop_gradient=True,[[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.]])
Tensor(shape=[5, 3], dtype=float32, place=Place(cpu), stop_gradient=True,[[2. , 3. , 4. ],[1. , 1. , 1. ],[8. , 9. , 10.],[1. , 1. , 1. ],[5. , 6. , 7. ]])

API 解析：index_add

查看前面的例子输出，可以看到，index_add就是把value的各个值，按照index里的值为索引，加入到源x里面去，比如

value = paddle.to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=paddle.float32)
index = paddle.to_tensor([0, 4, 2])

首先取出value[0] ,发现index[0]是 0,那么就把value[0] 跟x[0]相加

取出value[1] ,发现index[1] 是4,那么就把value[1] 跟x[4]相加

取出value[2] ,发现index[2] 是2,那么就把value[2] 跟x[2]相加

在飞桨官方没有index_add函数的时候，可以用python来实现，当然速度会慢很多：

def paddleindex_add(x, dim, index, source): # 飞桨的index_add'''
x = paddle.ones([5, 3])
t = paddle.to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=paddle.float32)
index = paddle.to_tensor([0, 4, 2])
# print(x)
with Benchmark("paddleindex_add"):x = paddleindex_add(x, 0, index, t)
print(x)'''for i in range(len(index)):x[index[i]] += source[i]return x

可以从赋值语句看到，就是从index里面取出值，然后x和source的相关值相加：x[index[i]] += source[i]

当然注释里面用了Benchmark函数，抄李沐老师的，源码如下

import time
class Timer:  #@save"""记录多次运行时间"""def __init__(self):self.times = []self.start()def start(self):"""启动计时器"""self.tik = time.time()def stop(self):"""停止计时器并将时间记录在列表中"""self.times.append(time.time() - self.tik)return self.times[-1]def avg(self):"""返回平均时间"""return sum(self.times) / len(self.times)def sum(self):"""返回时间总和"""return sum(self.times)def cumsum(self):"""返回累计时间"""return np.array(self.times).cumsum().tolist()class Benchmark:"""用于测量运行时间"""def __init__(self, description='Done'):self.description = descriptiondef __enter__(self):self.timer = Timer()return selfdef __exit__(self, *args):print(f'{self.description}: {self.timer.stop():.4f} sec')