飞桨Paddle API index_add 详解

index_add¶

paddle.index_add(x, index, axis, value, name=None)[源代码]¶

沿着指定轴 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 = description

    def __enter__(self):
        self.timer = Timer()
        return self

    def __exit__(self, *args):
        print(f'{self.description}: {self.timer.stop():.4f} sec')