Проблема преобразования программы Python в Cython для ускорения

Question

Я пытаюсь написать быстрый алгоритм на Cython. Алгоритм довольно прост и описан здесь (https://arxiv.org/pdf/1411.4357.pdf - параграф перед теоремой 8 на стр. 17). Идея относительно проста: предполагая, что данные редки, они могут быть заданы как входные данные в виде (i,j,A_ij) для матрицы A, которая имеет n строк и d столбцов. Теперь, чтобы воспользоваться преимуществом разреженности, нам нужны две функции h, которые отображают n строк в s блоков равномерно случайным образом (что является параметром для алгоритма) и знаковая функция s, которая равна ± 1 с равной вероятностью. Затем для каждой тройки (i,j,A_ij) алгоритм должен вывести (h(i),j, s(h(i))*A_ij) и вернуть их в виде массивов, которые будут переданы в качестве входных данных для другого разрежения.

Проблема в том, что я не могу получить ускорение. Это должно выполняться очень быстро и превосходить матричное умножение реализации S и умножения на A, как показано в приведенной выше ссылке.

Подход Python (примерно 11 мс):

    import numpy as np
    import numpy.random as npr
    from scipy.sparse import coo_matrix
    def countSketch_faster(input_matrix, sketch_size, seed=None):
        '''
        input_matrix: sparse - coo_matrix type
        sketch_size: int
        seed=None : random seed
        '''
        observed_rows = set([])
        sign_hashes = []
        hash_table = {}
        sketched_data = []
        hashed_rows = []


        for row_id, col_id, Aij in zip(input_matrix.row, input_matrix.col, input_matrix.data):
            if row_id in observed_rows:
                sketched_data.append(hash_table[row_id][1]*Aij)
                hashed_rows.append(hash_row_val)
            else:
                hash_row_val, sign_val = np.random.choice(sketch_size), np.random.choice([-1.0,1.0])
                hash_table[row_id] = (hash_row_val, sign_val) #hash-sign pair
                sketched_data.append(sign_val*Aij)
                hashed_rows.append(hash_row_val)
                observed_rows.add(row_id)


    hashed_rows = np.asarray(hashed_rows)
    sketched_data = np.asarray(sketched_data)
    row_hashes = np.asarray(row_hashes)

    S = coo_matrix((sketched_data, (hashed_rows, input_matrix.col))).tocsr()
    return S

Преобразование в Cython: Анализ аннотированных выходных данных показывает, что наиболее тяжелые строки на питоне - это те, которые называют numpy. Я попытался cdef все переменные и указать dtype для любого вызова numpy. Также я удалил команду zip и попытался сделать цикл более простым и более похожим на Си. Все это на самом деле оказало негативное влияние, и оно работает очень медленно, и я не совсем уверен, почему Кажется, что это довольно простой алгоритм для реализации, поэтому, если кто-то может помочь мне сократить время выполнения до чего-то очень маленького, я был бы чрезвычайно благодарен.

%%cython --annotate
cimport cython
import numpy.random as npr
import numpy as np
from scipy.sparse import coo_matrix

 def countSketch_cython(input_matrix, sketch_size, seed=None):
    '''
    input_matrix: sparse - coo_matrix type
    sketch_size: int
    seed=None : random seed
    '''

    cdef Py_ssize_t idx
    cdef int row_id
    cdef float data_id
    cdef float sketch_data_val
    cdef int hash_row_value
    cdef float sign_value



    hash_table = {}
    sketched_data = np.zeros_like(input_matrix.data,dtype=np.float64)
    hashed_rows = np.zeros_like(sketched_data,dtype=np.float64)
    observed_rows = set([])

    for idx in range(input_matrix.row.shape[0]):
        row_id = input_matrix.row[idx]
        data_id = input_matrix.data[idx]

        if row_id in observed_rows:
            hash_row_value = hash_table[row_id][0]
            sign_value = hash_table[row_id][1]
            sketched_data[row_id] = sign_value*data_id
            hashed_rows[idx] = hash_row_value
        else:
            hash_row_val = np.random.randint(low=0,high=sketch_size+1)
            sign_val = np.random.choice([-1.0,1.0])
            hash_table[row_id] = (hash_row_val, sign_val) #hash-sign pair
            sketched_data[idx] = sign_val*data_id
            hashed_rows[idx] = hash_row_value
            observed_rows.add(row_id)


    S = coo_matrix((sketched_data, (hashed_rows, input_matrix.col)))
return S`

ОБНОВЛЕНИЕ: мне удалось ускорить код, удалив некоторые из более медленных строк. Это были вызовы np.random, для которых генератор случайных чисел C работал быстрее, давая длину на входе, так что это не нужно было вычислять, и не преобразовывая в разреженную матрицу перед возвратом (для целей этого эксперимента я интересует, как быстро можно выполнить преобразование, а не детали, связанные с преобразованием для последующего использования).

%% cython --annotate Cimport Cython импортировать numpy.random as npr импортировать NumPy как NP из libc.stdlib cimport rand

#@cython.boundscheck(False)  # these don't contribute much in this 

пример #@cython.wraparound (False)
def countSketch (input_matrix, input_matrix_length, sketch_size, seed = None): «»» input_matrix: sparse - тип coo_matrix input_matrix_length - количество строк во входной матрице sketch_size: int seed = None: случайное семя '' '

    cdef Py_ssize_t idx
    cdef int row_id
    cdef float data_id
    cdef float sketch_data_val
    cdef int hash_row_value
    cdef float sign_value
    cdef int arr_lengths = input_matrix_length

    # These two lines are still annotated most boldly
    cdef double[:,:] sketched_data = 
np.zeros((arr_lengths,1),dtype=np.float64)
    cdef double[:,:] hashed_rows = np.zeros((arr_lengths,1))

    hash_table = {}

    observed_rows = set([])

    for idx in range(arr_lengths):
        row_id = input_matrix.row[idx]
        data_id = input_matrix.data[idx]

        if row_id in observed_rows:
            hash_row_value = hash_table[row_id][0]
            sign_value = hash_table[row_id][1]
            sketched_data[row_id] = sign_value*data_id
            hashed_rows[idx] = hash_row_value
        else:
            hash_row_val = rand()%(sketch_size) 
            #np.random.randint(low=0,high=sketch_size+1)
            sign_val = 2*rand()%(2) - 1
            #np.random.choice([-1.0,1.0])
            hash_table[row_id] = (hash_row_val, sign_val) #hash-sign pair
            sketched_data[idx] = sign_val*data_id
            hashed_rows[idx] = hash_row_value
            observed_rows.add(row_id)


    #S = coo_matrix((sketched_data, (hashed_rows, input_matrix.col)), dtype=np.float64) 
return hashed_rows, sketched_data

На случайной разреженной матрице A = scipy.sparse.random(1000, 50, density=0.1) теперь достигается 508 µs ± 17.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) с использованием %timeit countSketch(A, A.shape[0],100). Я должен представить, что все еще можно получить выгоду, так как я не знаю, правильно ли я настроил массивы.