Быстрый тест производительности Python для вычислительных задач

Мотивация

Совсем недавно вышла новая версия 0.34 библиотеки оптимизирующего JIT компилятора Numba для Python. И там ура! появилась долгожданная семантика аннотаций и набор методов для организации параллельных вычислений. За основу была взята технология Intel Parallel Accelerator.

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

Введение

В настоящее время Python очень активно используется в научных вычислениях, а в области Machine Learning вообще является практически одним из стандартов. Но если посмотреть чуть глубже, то почти везде Python используется как обертка над библиотеками более низкого уровня, написанных большей частью на C/C++. А можно ли на чистом Python писать на самом деле быстрый и параллельный код?

Рассмотрим совсем простую задачу. Пусть нам даны два набора N точек в трехмерном пространстве: p и q. Необходимо вычислить специальную матрицу на основе попарных расстояний между всеми точками:

.

Для всех тестов возьмем N = 5000. Время вычисления усредняется для 10 запусков.

Реализация на C++

Как точку отсчета возьмем следующую реализацию на C++:

void getR(std::vector<Point3D> & p, std::vector<Point3D> & q, Matrix & R)
{
    double rx, ry, rz;
    
    #pragma omp parallel for
    for (int i = 0; i < p.size(); i++) {
        for (int j = 0; j < q.size(); j++) {
	    rx = p[i].x - q[j].x;
            ry = p[i].y - q[j].y;
            rz = p[i].z - q[j].z;

            R(i, j) = 1 / (1 + sqrt(rx * rx + ry * ry + rz * rz));
        }
    }
}

Внешний цикл по точкам p выполняется параллельно с использованием технологии OpenMP.

Время выполнения: 44 мс.

Чистый Python

Начнем тест скорости с кода на чистом Python:

def get_R(p, q):
    R = np.empty((p.shape[0], q.shape[1]))

    for i in range(p.shape[0]):
        for j in range(q.shape[1]):
            rx = p[i, 0] - q[0, j]
            ry = p[i, 1] - q[1, j]
            rz = p[i, 2] - q[2, j]
            R[i, j] = 1 / (1 + math.sqrt(rx * rx + ry * ry + rz * rz))

    return R

Время выполнения: 52 861 мс, медленнее базовой реализации больше, чем в 1000 раз.

Python интерпретируемый язык, Python использует внутри себя GIL, что делает невозможным параллелизацию на уровне самого кода. Это все очень медленно. Теперь начнем это все ускорять.

Python + NumPy + SciPy

Проблему медленности Python для численных задач осознали очень давно. И ответом на эту проблему была библиотека NumPy. Идеология NumPy во многом близка MatLab, который является общепризнанным инструментом научных расчетов.

Мы прекращаем мыслить итерационно, мы начинаем мыслить матрицами и векторами как атомарными объектами для вычисления. А все операции с матрицам и векторами на нижнем уровне уже выполняются высокопроизводительными библиотеками линейной алгебры Intel MKL или ATLAS.

Реализация на NumPy выглядит так:

def get_R_numpy(p, q):
    Rx = p[:, 0:1] - q[0:1]
    Ry = p[:, 1:2] - q[1:2]
    Rz = p[:, 2:3] - q[2:3]

    R = 1 / (1 + np.sqrt(Rx * Rx + Ry * Ry + Rz * Rz))

    return R

В этой реализации вообще нет ни одного цикла!

Время выполнения: 839 мс, что медленнее базовой реализации где-то в 19 раз.

Более того, в NumPy и SciPy есть огромное количество встроенных функций. Реализация данной задачи на SciPy выглядит так:

def get_R_scipy(p, q):
    D = spd.cdist(p, q.T)
    R = 1 / (1 + D)
    return R

Время выполнения: 353 мс, что медленнее базовой реализации в 8 раз.

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

Но а как быть с параллельностью? Здесь она неявная. Мы надеемся, что на низком уровне все операции с матрицами и векторами реализованы эффективно и параллельно.

А что делать, если наш код не вписывается в линейную алгебру, или мы хотим явную параллелизацию?

Python + Cython

Тут приходит время Cython. Cython это специальный язык, который позволяет внутри обычного Python кода вставлять код на C-образом языке. Далее Cython преобразует это код в .c файлы, которые компилируется в python модули. Эти модули достаточно прозрачно можно вызывать в других частях Python кода. Реализация на Cython выглядит так:

@cython.wraparound(False)
@cython.nonecheck(False)
@cython.boundscheck(False)
def get_R_cython_mp(py_p, py_q):
    py_R = np.empty((py_p.shape[0], py_q.shape[1]))

    cdef int nP = py_p.shape[0]
    cdef int nQ = py_q.shape[1]
    cdef double[:, :] p = py_p
    cdef double[:, :] q = py_q
    cdef double[:, :] R = py_R

    cdef int i, j
    cdef double rx, ry, rz

    with nogil:
        for i in prange(nP):
            for j in xrange(nQ):
                rx = p[i, 0] - q[0, j]
                ry = p[i, 1] - q[1, j]
                rz = p[i, 2] - q[2, j]

                R[i, j] = 1 / (1 + sqrt(rx * rx + ry * ry + rz * rz))

    return py_R

Что тут происходит? На вход функция принимает python numpy объекты, далее они преобразуются в типизированные Cython С-структуры, а далее отключается gil и при помощи специальной конструкции 'prange' внешний цикл выполняется параллельно.

Время выполнения: 76 мс, что в 1.75 раз медленнее, чем базовая реализация.

Ну что, мы почти приблизились к базовой реализации, мы начали писать явный параллельный код. Но ценой этого стал менее читаемый код, мы ушли от чистого Python.

В целом, большинство численных расчетов так и пишутся. Большая часть на NumPy, а некоторые места критичные по быстродействию выносятся в отдельные модули и реализуются на cython.

Python + Numba

Мы проделали длинный путь. Мы начали с чистого Python, потом пошли по пути магии матричных вычислений, потом погрузились в специальных язык расширений. Пора вернуться обратно к тому, с чего мы начали. Итак, реализация на Python + Numba:

@jit(float64[:, :](float64[:, :], float64[:, :]), nopython=True, parallel=True)
def get_R_numba_mp(p, q):
    R = np.empty((p.shape[0], q.shape[1]))

    for i in prange(p.shape[0]):
        for j in range(q.shape[1]):
            rx = p[i, 0] - q[0, j]
            ry = p[i, 1] - q[1, j]
            rz = p[i, 2] - q[2, j]

            R[i, j] = 1 / (1 + math.sqrt(rx * rx + ry * ry + rz * rz))

    return R

Время выполнения: 46 мс, что практически совпадает с базовой реализацией!

И все, что нужно было сделать для этого с исходным медленным Python кодом:

• добавить аннотацию с указанием типов и параллельным типом выполнения
• заменить 'range' на 'prange' для параллельного выполнения внешнего цикла

Numba это очень интересный проект. Это оптимизирующий JIT компилятор для Python на основе LLVM. Он абсолютно прозрачен для использования, не нужно никаких отдельных модулей. Все что нужно, это добавить несколько аннотаций в критические по скорости методы.

В итоге, времена исполнения следующие:

Автор: alec_kalinin

Источник