Ожидается ли, что эффективный квадратный евклидов код расстояния numba будет медленнее, чем эффективный аналог numpy?

Question

Я изменяю наиболее эффективный код из ( Почему этот код Numba в 6 раз медленнее, чем код Numpy?), Чтобы он мог обрабатывать x1 (n, m)

@nb.njit(fastmath=True,parallel=True)
def euclidean_distance_square_numba_v5(x1, x2):
    res = np.empty((x1.shape[0], x2.shape[0]), dtype=x2.dtype)
    for a_idx in nb.prange(x1.shape[0]):
        for o_idx in range(x2.shape[0]):
            val = 0.
            for i_idx in range(x2.shape[1]):
                tmp = x1[a_idx, i_idx] - x2[o_idx, i_idx]
                val += tmp * tmp 
            res[a_idx, o_idx] = val 
    return res

Тем не менее, это еще не более эффективно, чем более эффективная версия numpy:

def euclidean_distance_square_einsum(x1, x2):
    return np.einsum('ij,ij->i', x1, x1)[:, np.newaxis] + np.einsum('ij,ij->i', x2, x2) - 2*np.dot(x1, x2.T)

С вводом как

a = np.zeros((1000000,512), dtype=np.float32)
b = np.zeros((100, 512), dtype=np.float32)

Время, которое я получил, составляет 2.4723422527313232 для кода numba и 0.8260958194732666 для кода numpy.

Answer 1

Да, это ожидается.

Первое, что вы должны знать: точка-продукт - это рабочая лошадка numpy-версии, здесь для немного меньших массивов:

>>> def only_dot(x1, x2):
        return - 2*np.dot(x1, x2.T)

>>> a = np.zeros((1000,512), dtype=np.float32)
>>> b = np.zeros((100, 512), dtype=np.float32)

>>> %timeit(euclidean_distance_square_einsum(a,b))
6.08 ms ± 312 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit(euclidean_only_dot(a,b))
5.25 ms ± 330 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

т.е. 85% времени уходит на это.

Когда вы смотрите на свой numba-код, это выглядит как несколько странная / необычная / более сложная версия умножения матрицы на матрицу - можно увидеть, например, те же три цикла.

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

Иногда, после всего веселья, нужно признать, что собственное "заново изобретенное колесо" не так хорошо, как современное колесо... но только тогда можно по-настоящему оценить его производительность.