Использование логической маски на большом массиве numpy очень медленно

Question

У меня проблема с производительностью при кодировании с python. скажем, у меня есть 2 очень больших массива (Nx2) строк, скажем, с N = 12 000 000, и две переменные label_a и label_b, которые также являются строками. Вот следующий код:

import numpy as np
import time

indices = np.array([np.random.choice(np.arange(5000).astype(str),size=10000000),np.random.choice(np.arange(5000).astype(str),size=10000000)]).T
costs = np.random.uniform(size=10000000)

label_a = '2'
label_b = '9'

t0 = time.time()    

costs = costs[(indices[:,0]!=label_a)*(indices[:,0]!=label_b)*(indices[:,1]!=label_a)*(indices[:,1]!=label_b)]
indices = indices[(indices[:,0]!=label_a)*(indices[:,0]!=label_b)*(indices[:,1]!=label_a)*(indices[:,1]!=label_b)]

t1 = time.time()
toseq = t1-t0
print(toseq)

указанный фрагмент кода занимает 3 секунды при каждом запуске. Я хотел бы добиться того же при одновременном снижении стоимости вычислений: я использую логическую маску для извлечения только строк в массивах затрат и индексов, где значения не являются label_a и label_b

Answer 1

Как указано в комментариях, вычисление значений индексов, которые вам нужны только один раз, и объединение их только один раз сэкономит время.

(я также изменил способ синхронизации, только для краткость - результаты совпадают)

import numpy as np
from timeit import timeit

r = 5000
n = 10000000

indices = np.array([
    np.random.choice(np.arange(r).astype(str), size=n),
    np.random.choice(np.arange(r).astype(str), size=n)
]).T
costs = np.random.uniform(size=n)

label_a = '2'
label_b = '9'

n_indices = np.array([
    np.random.choice(np.arange(r), size=n),
    np.random.choice(np.arange(r), size=n)
]).T


def run():
    global indices
    global costs

    _ = costs[(indices[:, 0] != label_a)*(indices[:, 0] != label_b) *
              (indices[:, 1] != label_a)*(indices[:, 1] != label_b)]
    _ = indices[(indices[:, 0] != label_a)*(indices[:, 0] != label_b) *
                (indices[:, 1] != label_a)*(indices[:, 1] != label_b)]


def run_faster():
    global indices
    global costs

    # only compute these only once
    not_a0 = indices[:, 0] != label_a
    not_b0 = indices[:, 0] != label_b
    not_a1 = indices[:, 1] != label_a
    not_b1 = indices[:, 1] != label_b
    _ = costs[not_a0 * not_b0 * not_a1 * not_b1]
    _ = indices[not_a0 * not_b0 * not_a1 * not_b1]


def run_even_faster():
    global indices
    global costs

    # also combine them only once
    cond = ((indices[:, 0] != label_a) * (indices[:, 0] != label_b) *
            (indices[:, 1] != label_a) * (indices[:, 1] != label_b))
    _ = costs[cond]
    _ = indices[cond]


def run_sep_mask():
    global indices
    global costs
    global cond

    # just the masking part of run_even_faster
    cond = ((indices[:, 0] != label_a) * (indices[:, 0] != label_b) *
            (indices[:, 1] != label_a) * (indices[:, 1] != label_b))


def run_sep_index():
    global indices
    global costs
    global cond

    # just the indexing part of run_even_faster
    _ = costs[cond]
    _ = indices[cond]


def run_even_faster_numerical():
    global indices
    global costs

    # use int values and n_indices instead of indices
    a = int(label_a)
    b = int(label_b)

    cond = ((n_indices[:, 0] != a) * (n_indices[:, 0] != b) *
            (n_indices[:, 1] != a) * (n_indices[:, 1] != b))
    _ = costs[cond]
    _ = indices[cond]


def run_all(funcs):
    for f in funcs:
        print('{:.4f} : {}()'.format(timeit(f, number=1), f.__name__))


run_all([run, run_faster, run_even_faster, run_sep_mask, run_sep_index, run_even_faster_numerical])

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

Этот прирост становится существенным, если вы начнете сравнивать более длинные метки - в конце концов, возможно, стоит даже конвертировать строки до чисел перед фильтрацией, если строки становятся достаточно длинными.

Вот мои результаты:

0.9711 : run()
0.7065 : run_faster()
0.6983 : run_even_faster()
0.2657 : run_sep_mask()
0.4174 : run_sep_index()
0.4536 : run_even_faster_numerical()

Две записи sep показывают, что индексирование примерно вдвое больше время, необходимое для создания маски для run_even_faster, так что вы можете ожидать такого улучшения только от еще большей настройки.

Однако они также показывают, что построение маски на основе целых чисел составляет менее 0,04 секунды. в дополнение к фактической индексации, по сравнению с примерно 0,26 секундами для построения маски на основе строк. Итак, у вас есть возможности для совершенствования.