Индекс по последним измерениям массива - PullRequest
Новая
       10

Индекс по последним измерениям массива

1 голос
/ 07 апреля 2020

Я бы хотел проиндексировать третье измерение массива, используя мой массив " indexer ", который содержит целочисленные индексы для этого третьего измерения.

Как мне избежать этого для l oop с помощью numpy методов? 

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 ничего не дал; Также не удалось найти точное соответствие для этого вопроса.

1 Ответ

2 голосов
/ 07 апреля 2020

Вот один из способов использования 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)
...