Ускорение расчета расстояния Numba - PullRequest
Новая
       5

Ускорение расчета расстояния Numba

0 голосов
/ 01 февраля 2020

Недавно я пытался вычислить расстояния до 2 верхних соседей в Python Numba следующим образом 

@jit(nopython=True)
def _latent_dim_kernel(data, pointers, indices, nrange, sampling_percentage = 1):

    pdists_t2 = np.zeros((nrange, 2))
    for a in range(nrange):

        rct = 0
        for b in range(nrange):
            if np.random.random() > 1- sampling_percentage:
                if a == b:
                    continue

                r1 = _get_sparse_row(a, data, pointers, indices)
                r2 = _get_sparse_row(b, data, pointers, indices)

                dist = np.linalg.norm(r2 - r1)

                if rct > 1:
                    if pdists_t2[a,0] > dist:
                        pdists_t2[a,0] = dist

                    elif pdists_t2[a,1] > dist:
                        pdists_t2[a,1] = dist
                else:
                    pdists_t2[a,rct] = dist

                rct += 1

    return pdists_t2

Данные, указатели и индексы: x.data, x.indptr, x. индексы матрицы CSR (scipy). Это работает нормально, однако, значительно медленнее, чем 

squareform(pdist(matrix)).sort(axis=1)[:,1:3]

Как я могу ускорить это без дополнительных затрат памяти?

Спасибо!

1 Ответ

0 голосов
/ 02 февраля 2020

Используйте попарные расстояния от sklearn

  • Поддерживаются попарные расстояния разреженных матриц (не требуется плотный временный массив)
  • В этом алгоритме используется перестановка алгебраических чисел c как в этом ответе
  • Это может быть намного быстрее при решении больших задач, таких как ваша (20k), поскольку большая часть вычислений выполняется в рамках высокооптимизированного матрично-матричного продукта.
  • Проверьте, является ли этот метод достаточно точным, он менее численно стабилен, чем «наивный» подход, используемый pdist

Пример 

import numpy as np
from scipy import sparse
from sklearn import metrics
from scipy.spatial import distance

matrix=sparse.random(1_000, 20_000, density=0.05, format='csr', dtype=np.float64)

%%timeit
dist_2=distance.squareform(distance.pdist(matrix.todense()))
dist_2.sort(axis=1)
dist_2=dist_2[:,1:3]
#10.1 s ± 23.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
dist=metrics.pairwise.euclidean_distances(matrix,squared=True)
dist.sort(axis=1)
dist=np.sqrt(dist[:,1:3])
#401 ms ± 13.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
...