数组最后维度的索引答案

【问题标题】：Index on last dimensions of array数组最后维度的索引
【发布时间】：2020-07-19 12:14:06
【问题描述】：

我想使用我的“indexer”数组对数组的第三维进行索引，该数组包含第三维的整数索引。

如何通过使用 numpy 方法避免执行此 for 循环？

import numpy as np

output = np.random.rand(3,12,40000) #batch x sequence x vocab_size, dtype float64
indexer = np.random.randint(0,40000, (3,12,52)) #batch x sequence x knowledgebase_size, dtype int64
values = np.random.rand(3,12,52) #batch x seq_len, knowledgebase_size, dtype float64

batch, sequence, kb = values.shape
for x in range(batch):
    for y in range(sequence):
        for z in range(kb):
            output[x,y,indexer[x,y,z]] += values[x,y,z]

查看 np 文档并没有产生任何结果；也找不到这个问题的完全匹配。

【问题讨论】：

标签： python arrays numpy multidimensional-array indexing

【解决方案1】：

这是一种使用advanced indexing 并使用np.ogrid 定义开放网格来索引output 并相应地添加values 的方法：

batch, sequence, kb = values.shape
i,j,_ = np.ogrid[:batch, :sequence, :kb]
output[i,j,indexer] += values

检查和时间安排 -

def adv_ix(out, values, indexer):
    batch, sequence, kb = values.shape
    i,j,k = np.ogrid[:batch, :sequence, :kb]
    out[i,j,indexer] += values
    return out

def current_app(out, values, indexer):
    batch, sequence, kb = values.shape
    for x in range(batch):
        for y in range(sequence):
            for z in range(kb):
                out[x,y,indexer[x,y,z]] += values[x,y,z]
    return out

np.allclose(adv_ix(output, values, indexer), current_app(output, values, indexer))
# True

%timeit adv_ix(output, values, indexer)
# 28.2 µs ± 341 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit current_app(output, values, indexer)
#1.49 ms ± 120 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

【讨论】：