Время прочтения: 6 мин.

При написании и отладки кода порой возникает вопрос «как же мне повысить скорость в Python», ведь как известно недостатками Python являются:

  1. Более низкая скорость работы;
  2. Более высокое потребление памяти написанных программ по сравнению с аналогичным кодом, написанным на компилируемых языках ( C или C++).

В этой статье рассмотрим библиотеку PyCUDA, как альтернативу CUDA для C/C++. Оценим её возможности и проведем сравнение производительности на конкретном примере, а именно реализуем алгоритм Харриса для детекции углов на изображении.

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

Алгоритм обнаружения углов Харриса (или Harris Corner Detector) – это оператор детектирования углов, который обычно используется в алгоритмах компьютерного зрения. Впервые этот алгоритм был представлен К. Харрисом и М. Стивенсом в 1988 году и с того времени он был улучшен и принят в различные алгоритмы для предварительной обработки изображений.

Также существует множество вариаций и дополнений к алгоритму Харриса, таких как алгоритм обнаружения углов Ши-Томаси, детектор Трайковича-Хедли и оператор Моравека. Большинство из них выполняют функцию оценки яркости с некоторым окном и функцией активации. Такие алгоритмы требуют предварительной обработки изображения, например, размытия по Гауссу. Тем не менее ни один алгоритм не учитывает особенности CUDA, но все они могут быть легко распараллелены на графических процессорах, поскольку каждый пиксель обрабатывается отдельно, однако ему нужны значения из соседних пикселей.

В данной статье будем использовать Google Colaboratory, так как в рамках статьи нам хочется проверить скорость работы на CPU и GPU соответственно. Проверим, какие CPU и GPU нам выдадут в Google Colaboratory:

Вначале рассмотрим реализацию детектора Харриса на CPU. Код для алгоритма Харриса будет выглядеть следующим образом:

start = time.time()

filename = '/content/drive/MyDrive/image/art_4.png'
img = cv2.imread(filename)
gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

gray = np.float32(gray)
dst = cv2.cornerHarris(gray,2,3,0.04)

#result is dilated for marking the corners, not important
dst = cv2.dilate(dst,None)

# Threshold for an optimal value, it may vary depending on the image.
img[dst>0.01*dst.max()]=[0,0,255]

print('Результат алгоритма:')
cv2_imshow(img)
if cv2.waitKey(0) & 0xff == 27:
    cv2.destroyAllWindows()

end  = time.time() - start
print('Время вычисления на CPU: {:2.4f}'.format(end))

Выполнив данный блок кода, получаем результат алгоритма и время выполнения этого алгоритма на CPU:

Время вычисления на CPU: 0.4256 сек.

Заметим, что скорость выполнения составила почти полторы секунды, что для кода, использующего OpenCV неплохо. Но возможно ли сделать манипуляции с кодом, чтобы добиться ускорения?

Рассмотрим эту возможность на примере реализации все того же кода, но уже с применением библиотеки PyCUDA.

Реализуем алгоритм Харриса для детектирования углов и распараллелим процессы на GPU, чтобы повысить производительность кода.

def Harris_GPU(img, k, thresh):

    height = img.shape[0]
    width = img.shape[1]

    vector_size = img.shape[0] * img.shape[1]
    corner_list = []
    offset = 2
    # to fit still in a 32-bit integer
    thresh = int(thresh/10)

    # function template
    func_mod_template = Template("""
    #include<stdio.h>
    #define INDEX(a, b) a*${HEIGHT}+b

    __global__ void corners(
        float *dest,
        float *ixx,
        float *ixy,
        float *iyy,
        int offset,
        float k,
        int threshold) {

        unsigned int idx = threadIdx.x + threadIdx.y*blockDim.y +
                            (blockIdx.x*(blockDim.x*blockDim.y));

        unsigned int a = idx/${HEIGHT};
        unsigned int b = idx%${HEIGHT};

        float sxx = 0;
        float sxy = 0;
        float syy = 0;
        float det = 0;
        float trace = 0;
        float r = 0;

        if ((a >= offset) & (a <= (${WIDTH}-offset - 1)) &
            (b >= offset) & (b <= (${HEIGHT}-offset - 1))) {
            for (int bi = b - offset; bi < b + offset + 1; ++bi) {
                for (int ai = a - offset; ai < a + offset + 1; ++ai) {
                    sxx = sxx + ixx[INDEX(ai, bi)];
                    sxy = sxy + ixy[INDEX(ai, bi)];
                    syy = syy + iyy[INDEX(ai, bi)];
                }
            }
            det = sxx*syy - sxy*sxy;
            trace = sxx + syy;
            r = det - k*(trace*trace);
            if ((r/10) > threshold)
                dest[INDEX(a, b)] = r;
        }
    }
    """)

    func_mod = SourceModule(func_mod_template.substitute(HEIGHT=height,
                                                         WIDTH=width))
    pycuda_corners = func_mod.get_function("corners")

    # Find x and y derivatives
    dy, dx = np.gradient(img)
    Ixx = dx**2
    Ixy = dy*dx
    Iyy = dy**2

    ixx = Ixx.reshape(vector_size, order='F')
    ixy = Ixy.reshape(vector_size, order='F')
    iyy = Iyy.reshape(vector_size, order='F')
    dest_r = np.zeros_like(ixx)

    # start timer
    start = timeit.default_timer()

    pycuda_corners(drv.Out(dest_r),
                drv.In(ixx),
                drv.In(ixy),
                drv.In(iyy),
                np.uint32(offset),
                np.float32(k),
                np.uint32(thresh),
                block=(32, 32, 1),  
                grid=(height*width//1024, 1, 1))  

    # calculate used time
    execution_time = timeit.default_timer() - start

    # extract the corners
    r = np.reshape(dest_r, (height, width), order='F')
    corners = np.where(r > 0)
    for i, j in zip(corners[0], corners[1]):
        corner_list.append([j, i, r[i, j]])

    return corner_list, execution_time

Отметим, что часть кода заимствует синтаксис C/С++, что позволяет ускорить работу нашего кода, написанного на Python.

Получаем результат и время работы кода на GPU:

Время вычисления на GPU: 0.2576 сек.

Хочется отметить, что реализация на GPU детектирует далеко не все «особенные» точки и весь процесс зацикливается в основном на левой части картинки, когда в свою очередь реализация на CPU детектирует большее количество углов на изображении. Тем не менее скорость вычисления на GPU быстрее, а именно почти в два раза, чем на CPU.

В качестве вывода отметим то, что для реализации крупных и «тяжелых» проектов или задач, где скорость важна и необходима, использовать библиотеку PyCUDA будет плюсом (опять-таки если знаете C/C++).

Добавлю, что PyCUDA является библиотекой с открытым исходным кодом и лицензией MIT. Документация понятная и ясная для любого пользователя, к тому же имеет множество находящихся в открытом доступе примеров кода, которые могут оказать содействие и поддержку. Главная цель использования PyCUDA состоит в том, чтобы позволить пользователю применять необходимые шаблоны CUDA с минимумом абстракций Python.