Оптимизация набора кортежей целых чисел с помощью Numba?

Я учусь использовать Numba (хотя я уже достаточно знаком с Cython). Как мне ускорить этот код? Обратите внимание, что функция возвращает список наборов из двух кортежей целых чисел. Я использую блокнот IPython. Я бы предпочел Numba Cython.

@autojit
def generateadj(width,height):
    adj = {}
    for y in range(height):
        for x in range(width):
            s = set()
            if x>0:
                s.add((x-1,y))
            if x<width-1:
                s.add((x+1,y))
            if y>0:
                s.add((x,y-1))
            if y<height-1:
                s.add((x,y+1))
            adj[x,y] = s
    return adj

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

%%cython
import numpy as np

def generateadj(int width, int height):
    cdef int[:,:,:,:] adj = np.zeros((width,height,4,2), np.int32)
    cdef int count

    for y in range(height):
        for x in range(width):
            count = 0
            if x>0:
                adj[x,y,count,0] = x-1
                adj[x,y,count,1] = y
                count += 1
            if x<width-1:
                adj[x,y,count,0] = x+1
                adj[x,y,count,1] = y
                count += 1
            if y>0:
                adj[x,y,count,0] = x
                adj[x,y,count,1] = y-1
                count += 1
            if y<height-1:
                adj[x,y,count,0] = x
                adj[x,y,count,1] = y+1
                count += 1
            for i in range(count,4):
                adj[x,y,i] = adj[x,y,0]
    return adj

python numba cython
person ArekBulski    schedule 11.01.2015    source источник

Ответы (1)


arrow_upward
5
arrow_downward

Хотя numba поддерживает такие структуры данных Python, как dicts и sets, он делает это в режиме объекта. Из глоссария numba объектный режим определяется как:

Режим компиляции Numba, который генерирует код, обрабатывающий все значения как объекты Python, и использует Python C API для выполнения всех операций с этими объектами. Код, скомпилированный в объектном режиме, часто будет работать не быстрее, чем код, интерпретируемый Python, если только компилятор Numba не может воспользоваться преимуществами цикла.

Поэтому при написании кода numba вам нужно придерживаться встроенных типов данных, таких как массивы. Вот код, который делает именно это:

@jit
def gen_adj_loop(width, height, adj):
    i = 0
    for x in range(width):
        for y in range(height):
            if x > 0:
                adj[i,0] = x
                adj[i,1] = y
                adj[i,2] = x - 1
                adj[i,3] = y
                i += 1

            if x < width - 1:
                adj[i,0] = x
                adj[i,1] = y
                adj[i,2] = x + 1
                adj[i,3] = y
                i += 1

            if y > 0:
                adj[i,0] = x
                adj[i,1] = y
                adj[i,2] = x
                adj[i,3] = y - 1
                i += 1

            if y < height - 1:
                adj[i,0] = x
                adj[i,1] = y
                adj[i,2] = x
                adj[i,3] = y + 1
                i += 1
    return

Это принимает массив adj. Каждая строка имеет вид x y adj_x adj_y. Итак, для пикселя (3,4) у нас будет четыре строки:

3 4 2 4
3 4 4 4
3 4 3 3
3 4 3 5

Мы можем обернуть вышеуказанную функцию в другую:

@jit
def gen_adj(width, height):
    # each pixel has four neighbors, but some of these neighbors are
    # off the grid -- 2*width + 2*height of them to be exact
    n_entries = width*height*4 - 2*width - 2*height
    adj = np.zeros((n_entries, 4), dtype=int)
    gen_adj_loop(width, height, adj)

Эта функция очень быстрая, но неполная. Мы должны преобразовать adj в словарь формы в вашем вопросе. Проблема в том, что это очень медленный процесс. Мы должны перебрать массив adj и добавить каждую запись в словарь Python. Это не может быть изменено numba.

Итак, суть такова: требование, чтобы результатом был словарь кортежей, действительно ограничивает возможности оптимизации этого кода.

person jme    schedule 11.01.2015
comment
Обратите внимание, что numba добавил поддержку многочисленных функций с даты этого ответа. - person songololo; 31.12.2019