NumPy — делаем быстрее

Чтобы ускорить выполнение кода на интерпретируемом языке программирования, следуют известным советам по избеганию использования циклов, своевременному выделению памяти, применению параллелизации и встроенных средств линейной алгебры. Для М-языка (Matlab) все советы можно найти в гл.15 уже старенькой книги автора. В принципе, они все годятся и для библиотеки NumPy языка Python, поскольку она практически копирует функционал Матлаба. Но кроме самих советов есть некоторые тонкости их использования. Сейчас на одном конкретном примере покажем, как можно ускорить выполнение кода.

Пример взят из соревнования BlackBoxChallenge (идёт в настоящий момент), в котором надо написать агента, играющего в некоторую игру (правила неизвестны и агент должен сам научиться в неё играть). Не будем вдаваться в детали, рассмотрим пример агента, который предоставили организаторы (см. листинг). В каждый момент времени известно состояние игры – state – 36-мерный вещественный вектор, надо сделать одно из 4х возможный действий. Простейший алгоритм использует идею линейной регрессии: каждому действию поставить в соответствие некоторую линейную функцию, в каждый момент времени вычислить значения этих 4х линейных функций и совершить действие, которое соответствует максимальному значению. Как настроить параметры линейных функций – это задача машинного обучения, а мы сейчас поговорим о том, как эффективнее реализовать этот подход.

Изначально коэффициенты лежат в 4×36-матрице coefs, она «растаскивается» на матрицу самих коэффициентов и вектор свободных членов, первая матрица умножается на вектор состояния state, а затем прибавляется вектор свободных членов, максимум выбирается в цикле (см. рис.). Покажем, как средствами NumPy ускорить этот код и сделать его компактным.

Рис. Схематическое изображение используемого подхода.

Первый вариант такой (см. листинг). Коэффициенты храним в матрице 4×37, «растаскивание» происходит перед самим умножением без использования лишних переменных (такой подход – хранение всех параметров в одном месте  – существенно упрощает машинное обучение). Максимум выбирается не в цикле, а с помощью функции argmax. Вместо вызова двух функций – две строчки кода (хотя, организаторы конечно специально сделали функции – для объяснения логики работы агента).

Второй подход – никакого растаскивания. Просто пополняем вектор состояния фиктивным константным элементом (признаком) с помощью функции append (типичный приём у программистов нейронных сетей).

Третий подход – такой же как и второй, но вместо вызова функции append на каждом шаге, мы просто делаем шаблон вектора состояний state и заполняем его динамическую часть.

На диаграмме показано, сколько по времени заняло исполнение всех представленных фрагментов кода. В результате получилось

• Почти в 2.5 раза ускорить работу агента
• Все параметры хранить в одной переменной (типа np.array, что очень удобно)
• Сделать код компактным

Автор не сравнивал представленные фрагменты кода с эквивалентной Cython-реализацией. В принципе, это интересно: можно ли существенно быстрее реализовать линейные алгоритмы с помощью других библиотек?