Эффективно рассчитайте сумму секционирования с помощью CuPy или NumPy - PullRequest
Новая
       79

Эффективно рассчитайте сумму секционирования с помощью CuPy или NumPy

1 голос
/ 10 июля 2020

У меня есть очень длинный массив * длины L (назовем его values), который я хочу суммировать, и отсортированный 1D-массив той же длины L, содержащий N целые числа, с которыми чтобы разбить исходный массив - назовем этот массив labels.

То, что я сейчас делаю, это (module является cupy или numpy): 

result = module.empty(N)
for i in range(N):
    result[i] = values[labels == i].sum()

Но это не может быть наиболее эффективным способом сделать это (должно быть возможно избавиться от for l oop, но как?). Поскольку labels отсортировано, я мог легко определить точки останова и использовать эти индексы в качестве точек начала / остановки, но я не понимаю, как это решает проблему for l oop.

Примечание что я бы хотел избежать создания массива размером N x L по пути, если это возможно, поскольку L очень большой.

Я работаю в cupy, но с любым numpy решение тоже приветствуется и, вероятно, может быть перенесено. В случае с Cupy, похоже, это будет случай ReductionKernel, но я не совсем понимаю, как это сделать.

* в моем случае, values - 1D, но я предполагаю, что решение не будет зависеть от этого

Ответы [ 2 ]

2 голосов
/ 11 июля 2020

Вы описываете групповое агрегирование по сумме. Вы можете написать для этого CuPy RawKernel, но было бы намного проще использовать существующие агрегаты groupby, реализованные в cuDF , библиотеке фреймов данных GPU. Они могут взаимодействовать, не требуя от вас копирования данных. Если вы вызовете .values в результирующей серии cuDF, он даст вам массив CuPy.

Если вы вернетесь к процессору, вы можете сделать то же самое с pandas. 

import cupy as cp
import pandas as pd

N = 100
values = cp.random.randint(0, N, 1000)
labels = cp.sort(cp.random.randint(0, N, 1000))

L = len(values)
result = cp.empty(L)
for i in range(N):
    result[i] = values[labels == i].sum()
    
result[:5]
array([547., 454., 402., 601., 668.])

import cudf

df = cudf.DataFrame({"values": values, "labels": labels})
df.groupby(["labels"])["values"].sum().values[:5]
array([547, 454, 402, 601, 668])
1 голос
/ 12 июля 2020

Вот решение, которое вместо массива N x L использует массив N x <max partition size in labels> (который не должен быть большим, если несоответствие между различными разделами не слишком велико):

  1. Измените размер массива в двумерный массив с разделами в каждой строке. Поскольку длина строки равна размеру максимального раздела, заполните недоступные значения нулями (так как это не влияет на сумму). Здесь используется решение @Divakar, здесь . 
def jagged_to_regular(a, parts):
    lens = np.ediff1d(parts,to_begin=parts[0])
    mask = lens[:,None]>np.arange(lens.max())
    out = np.zeros(mask.shape, dtype=a.dtype)
    out[mask] = a
    return out

parts_stack = jagged_to_regular(values, labels)
Сумма по оси 1: 
result = np.sum(parts_stack, axis = 1)

Если вам нужна реализация CuPy, прямой альтернативы CuPy от до numpy.ediff1d в jagged_to_regular не существует. В этом случае вы можете заменить оператор на numpy.diff, например: 

lens = np.insert(np.diff(parts), 0, parts[0])

, а затем продолжить использовать CuPy в качестве замены numpy.

...