Алгоритм тензордота, реализованный в Numba, намного медленнее, чем у Numpy. - PullRequest
Новая
       9

Алгоритм тензордота, реализованный в Numba, намного медленнее, чем у Numpy.

0 голосов
/ 11 февраля 2019

Я пытаюсь расширить numy "tenordot" так, чтобы такие вещи, как: K_ijklm = A_ki * B_jml можно было записать понятным образом, например: K = mytensordot(A,B,[2,0],[1,4,3])

Насколько я понимаю, tenordot numpy (с необязательным аргументом0) смог бы сделать что-то вроде этого: K_kijml = A_ki * B_jml, т.е. сохранить порядок индексов.Поэтому мне нужно было бы сделать ряд np.swapaxes(), чтобы получить матрицу `K_ijklm ', которая в сложном случае может быть легким источником ошибок (потенциально очень сложным для отладки).

Проблема в том,что моя реализация идет медленно (в 10 раз медленнее, чем тензордо [РЕДАКТИРОВАТЬ: на самом деле это НАМНОГО медленнее, чем это)), даже при использовании numba.Мне было интересно, сможет ли кто-нибудь понять, что можно сделать для улучшения производительности моего алгоритма.

MWE 

import numpy as np
import numba as nb
import itertools
import timeit

@nb.jit()
def myproduct(dimN):
    N=np.prod(dimN)
    L=len(dimN)
    Product=np.zeros((N,L),dtype=np.int32)
    rn=0
    for n in range(1,N):
        for l in range(L):
            if l==0:
                rn=1
            v=Product[n-1,L-1-l]+rn
            rn = 0
            if v == dimN[L-1-l]:
                v = 0
                rn = 1
            Product[n,L-1-l]=v
    return Product

@nb.jit()
def mytensordot(A,B,iA,iB):
    iA,iB = np.array(iA,dtype=np.int32),np.array(iB,dtype=np.int32)
    dimA,dimB = A.shape,B.shape
    NdimA,NdimB=len(dimA),len(dimB)

    if len(iA) != NdimA: raise ValueError("iA must be same size as dim A")
    if len(iB) != NdimB: raise ValueError("iB must be same size as dim B")

    NdimN = NdimA + NdimB
    dimN=np.zeros(NdimN,dtype=np.int32)
    dimN[iA]=dimA
    dimN[iB]=dimB
    Out=np.zeros(dimN)
    indexes = myproduct(dimN)

    for nidxs in indexes:
        idxA = tuple(nidxs[iA])
        idxB = tuple(nidxs[iB])
        v=A[(idxA)]*B[(idxB)]
        Out[tuple(nidxs)]=v
    return Out



A=np.random.random((4,5,3))
B=np.random.random((6,4))

def runmytdot():
    return mytensordot(A,B,[0,2,3],[1,4])
def runtensdot():
    return np.tensordot(A,B,0).swapaxes(1,3).swapaxes(2,3)


print(np.all(runmytdot()==runtensdot()))
print(timeit.timeit(runmytdot,number=100))
print(timeit.timeit(runtensdot,number=100))

Результат: 

True
1.4962144780438393
0.003484356915578246

Ответы [ 2 ]

0 голосов
/ 11 февраля 2019

tensordot со значениями скалярных осей могут быть неясными.Я исследовал это в

Как пошагово работает функция numpy.tensordot?

Там я понял, что np.tensordot(A, B, axes=0) эквивалентно, используя axes=[[], []]. 

In [757]: A=np.random.random((4,5,3))
     ...: B=np.random.random((6,4))

In [758]: np.tensordot(A,B,0).shape
Out[758]: (4, 5, 3, 6, 4)
In [759]: np.tensordot(A,B,[[],[]]).shape
Out[759]: (4, 5, 3, 6, 4)

Это, в свою очередь, эквивалентно вызову dot с новым размерным суммой продуктов размера 1: 

In [762]: np.dot(A[...,None],B[...,None,:]).shape
Out[762]: (4, 5, 3, 6, 4)

(4,5,3,1) * (6,1,4)   # the 1 is the last of A and 2nd to the last of B

dot быстро, с использованием BLAS (или эквивалентного) код.Смена осей и изменение формы также происходит относительно быстро.

einsum дает нам большой контроль над осями

, повторяя перечисленные выше продукты: 

In [768]: np.einsum('jml,ki->jmlki',A,B).shape
Out[768]: (4, 5, 3, 6, 4)

и с заменой:

In [769]: np.einsum('jml,ki->ijklm',A,B).shape
Out[769]: (4, 4, 6, 3, 5)

Незначительный момент - двойной своп можно записать как одну транспонированную: 

.swapaxes(1,3).swapaxes(2,3)
.transpose(0,3,1,2,4)
0 голосов
/ 11 февраля 2019

Вы столкнулись с известной проблемой .numpy.zeros требует кортежа при создании многомерного массива.Если вы передаете что-то, кроме кортежа, это иногда работает, но это только потому, что numpy умен, чтобы сначала преобразовать объект в кортеж.

Проблема в том, что numba в настоящее время не поддерживает преобразование произвольных итераций в кортежи .Так что эта строка не работает, когда вы пытаетесь скомпилировать ее в режиме nopython=True.(Несколько других тоже терпят неудачу, но это первое.) 

Out=np.zeros(dimN)

Теоретически вы можете вызвать np.prod(dimN), создать плоский массив нулей и изменить его, но затем вы столкнетесь ста же проблема: метод reshape для массивов numpy требует кортежа!

Это довольно неприятная проблема с numba - я не сталкивался с этим раньше.Я действительно сомневаюсь, что решение, которое я нашел, является правильным, но это рабочее решение, которое позволяет нам скомпилировать версию в режиме nopython=True.

Основная идея состоит в том, чтобы избежать использования кортежей для индексации путем непосредственной реализации индексатора, который следует за strides массива: 

@nb.jit(nopython=True)
def index_arr(a, ix_arr):
    strides = np.array(a.strides) / a.itemsize
    ix = int((ix_arr * strides).sum())
    return a.ravel()[ix]

@nb.jit(nopython=True)
def index_set_arr(a, ix_arr, val):
    strides = np.array(a.strides) / a.itemsize
    ix = int((ix_arr * strides).sum())
    a.ravel()[ix] = val

Это позволяет нам получать и устанавливать значения без необходимостикортеж

Мы также можем избежать использования reshape, передав выходной буфер в функцию joted и обернув эту функцию в помощник: 

@nb.jit()  # We can't use nopython mode here...
def mytensordot(A, B, iA, iB):
    iA, iB = np.array(iA, dtype=np.int32), np.array(iB, dtype=np.int32)
    dimA, dimB = A.shape, B.shape
    NdimA, NdimB = len(dimA), len(dimB)

    if len(iA) != NdimA:
        raise ValueError("iA must be same size as dim A")
    if len(iB) != NdimB:
        raise ValueError("iB must be same size as dim B")

    NdimN = NdimA + NdimB
    dimN = np.zeros(NdimN, dtype=np.int32)
    dimN[iA] = dimA
    dimN[iB] = dimB
    Out = np.zeros(dimN)
    return mytensordot_jit(A, B, iA, iB, dimN, Out)

Поскольку помощник не содержит циклов, он добавляетнекоторые накладные расходы, но накладные расходы довольно тривиальны.Вот последняя заключенная в кавычки функция: 

@nb.jit(nopython=True)
def mytensordot_jit(A, B, iA, iB, dimN, Out):
    for i in range(np.prod(dimN)):
        nidxs = int_to_idx(i, dimN)
        a = index_arr(A, nidxs[iA])
        b = index_arr(B, nidxs[iB])
        index_set_arr(Out, nidxs, a * b)
    return Out

К сожалению, это не приводит к такому ускорению, как нам хотелось бы.На массивах это примерно в 5 раз медленнее, чем tensordot;на больших массивах он все еще в 50 раз медленнее.(Но, по крайней мере, это не в 1000 раз медленнее!) В ретроспективе это не слишком удивительно, поскольку dot и tensordot оба используют BLAS под капотом, так как @hpaulj напоминает нам .

Закончив этот код, я увидел, что einsum решил вашу настоящую проблему - приятно!

Но основная проблема, на которую указывает ваш первоначальный вопрос - что индексация с кортежами произвольной длины невозможна в объединенном коде - все еще вызывает разочарование.Надеюсь, это будет полезно кому-то еще!

...