Использование Numba с Numpy для ускорения Python
Python. Благословенен ты среди языков программирования, прости меня за то, что я грешник. В момент своего высокомерия я усомнился в твоей скорости, не понимая, что ты дал своим ягнятам тупых, как спасителя в тяжелые времена.
Для некоторых людей может стать откровением, что Python работает медленно. Для них не торопитесь, чтобы преодолеть пять стадий горя. Теперь, когда это изучено, давайте научимся numba выжимать производительность из кода Python.
Инициализация
Чтобы установить numba, выполните следующую команду. Вы можете заменить pip на conda или pip3 в зависимости от используемого менеджера пакетов.
$ pip install numba
Для создания теста нам нужно создать среду. Я использую Python3.8. Моя ОС - ubuntu 20.04. Мы будем использовать массив из 10 миллионов чисел с плавающей запятой. Мы постараемся получить квадратный корень из всех элементов и сохранить его в отдельном массиве. Я беру в среднем 100 тестов для создания теста. Скорость в вашей реализации может изменяться до + -10%. Давайте сначала создадим массив.
import numpy as np size = int(5e7) array = np.arange(size).astype(np.float32)
Теперь мы будем использовать обычную операцию Python для создания контрольной точки. Затем мы попробуем оптимизировать производительность с помощью numba. Давайте начнем с простого и добавим сложности по мере продвижения.
Обычный цикл: {9,29 секунды}
def function_with_normal_loop(array):
out = []
for i in array:
out.append(sqrt(i))
return out
Мы перебираем массив, извлекаем квадратный корень и добавляем его в массив. Время, необходимое для выполнения цикла, составляет 9.2 секунд. Причина, по которой он медленный, заключается в том, что python - это интерпретируемый язык. Таким образом, количество строк в Python может иметь значение.
Понимание списка: {7,34 секунды}
def function_with_list_comprehension(array):
return [sqrt(x) for x in array]
Мы используем питонический способ перебора массива. Функция занимает 7.3 секунд. Это улучшение на 2 секунды и на 3 строки меньше кода. Так впечатляет.
Карта: {5,9 секунды}
def function_with_map(array):
return list(map(sqrt, array))
map() - встроенная функция. Он проходит через итератор и возвращает генератор, который необходимо преобразовать в список. Причина, по которой карта работает быстрее, чем понимание списка, заключается в том, что карта реализована на C.
Теперь, когда мы установили тест для нормальной функции Python, давайте попробуем использовать numba.
JIT
Numba поставляется с Just-in-time compiler (JIT). Его задача - преобразовать интерпретируемый код в байт-код, чтобы при следующем вызове функции не нужно было повторно интерпретировать ее во время выполнения. Чтобы использовать его, мы добавляем decorator@jit(nopython=True. Для его импорта мы используем:
from numba import jit
Обычный цикл с JIT: {2,24 секунды}
@jit(nopython=True)
def function_with_normal_loop(array):
out = []
for i in array:
out.append(sqrt(i))
return out
Мы видим заметное увеличение скорости примерно с 9 секунд до 2 секунд. И это просто добавление заголовка.
Понимание списка с помощью JIT: {2,34 секунды}
@jit(nopython=True)
def function_with_list_comprehension(array):
return [sqrt(x) for x in array]
Это было антиклиматически. Скорость между пониманием списка без JJIT и с JIT увеличилась, но не настолько, насколько мы могли бы видеть с обычным циклом. Кажется, мы достигли плато при увеличении скорости с помощью Numba.
Карта с JIT: {4,14 секунды}
@jit(nopython=True)
def function_with_map(array):
return list(map(sqrt, array))
Вот это просто грустно. Numba работает только с примитивными типами данных и обычными циклами. Но для уловок, которые мы используем, чтобы сделать Python быстрее в обычном цикле, когда мы используем JIT.
Использование JIT в первый раз связано с накладными расходами. Поэтому мы должны использовать его только для функций, которые вызываются повторно. И мы должны протестировать код, чтобы увидеть, насколько улучшилось использование numba.
Векторизовать
Векторизация выполняет операцию, которая может быть выполнена в отдельном элементе, и выполняет ее для всего массива. Для импорта нам нужно сделать
from numba import vectorize
Векторизация: {0,04 секунды}
@vectorize
def function_with_vectorize(elem):
return sqrt(elem)
В первый раз мы запускаем код. Это займет 0.116 секунд, потому что он пытается найти тип данных для функции. После этого каждый раз, когда мы вызываем функцию, код будет запускать 0.047 seconds. Это разумное увеличение производительности. Посмотрим, сможем ли мы сохранить производительность для первой итерации.
@vectorize([‘float32(float32)’])
def function_with_vectorize(elem):
return sqrt(elem)
Функция такая же, как и функция выше, за исключением того, что мы указали numba, что тип ввода и тип вывода как массив float32. Теперь с самого начала мы получим ту же производительность.
CUDA
Numba также позволяет нам использовать графический процессор вместо процессора. GPU полезны для массового параллельного программирования. И может значительно улучшить скорость вашей программы. Для его импорта мы будем использовать:
from numba import cuda
JIT с CUDA: {0,08 секунды}
@cuda.jit
def normal_function(array, out):
idx = cuda.grid(1)
out[idx] = sqrt(array[idx])
d_a = cuda.to_device(array)
d_out = cuda.device_array_like(d_a)
blocks_per_grid = 32
threads_per_block = 128
normal_function[blocks_per_grid, threads_per_block](d_a, d_out)
print(d_out.copy_to_host())
Устройство - это графический процессор, а хост - центральный процессор. У них есть собственная DRAM. Вот почему данные копируются с хоста на устройство и обратно. Мы также должны определить количество потоков и блоков для выполнения. На данный момент информации может быть слишком много, о ней мы поговорим в следующей статье. Для упрощения рассмотрим другой способ использования графического процессора.
Векторизация с помощью CUDA: {0,08 секунды}
Мы можем использовать jit для CUDA, чтобы добиться той же производительности. Код такой же, как и у обычной векторной операции, которую мы использовали для@vectorize. Чтобы использовать его, мы просто добавляем target=’cuda’ к тому же коду.
@vectorize(['float32(float32)'], target='cuda')
def vectorize_with_cuda(elem, ):
return sqrt(elem)
GPU работает быстро при параллельной обработке, но при запуске ядра возникают накладные расходы. И данные должны быть переданы в GPU от CPU, и после выполнения их нужно скопировать обратно. Это общение - узкое место. Более подробная информация находится здесь. Также одномерный массив не является способом отображения мощности графического процессора. В следующей статье мы обсудим, как использовать CUDA для программирования GPU на Python, чтобы еще больше повысить скорость.
Заключение
Только с этим мы улучшили производительность более чем в 200 раз. Это не лучшая производительность, которую мы можем получить, но это хороший момент для начала. Надеюсь, вам эта статья показалась интересной и полезной.
Больше контента на plainenglish.io
