Простая оптимизация производительности кода Python.

Python имеет репутацию производителя медленно работающего кода. На первый взгляд это может быть правдой, но при ближайшем рассмотрении оказывается несправедливым суждением. На самом деле, в любом программном проекте производительность — это проблема лишь нескольких узких мест, независимо от того, используем ли мы Python, C или Java. Конечно, эти узкие места должны быть устранены. Прелесть Python в том, что он идеально взаимодействует с низкоуровневыми языками, такими как C, что позволяет писать высокооптимизированный код для нескольких узких мест. Так что на самом деле нет необходимости отказываться от Python из-за предполагаемых соображений производительности. Обычно нам даже не нужно замарать руки самим низкоуровневым кодом, потому что другие люди уже сделали эту работу, и мы можем остаться в нашем мире Python и наслаждаться отличной производительностью. Или вы думаете, что Python был бы первоклассным гражданином в мире глубокого обучения с тяжелыми вычислениями? Позвольте мне показать вам несколько простых способов значительно ускорить код Python.

Множество Мандельброта

В качестве примера оптимизации Python мы реализуем множество Мандельброта, которое является самым известным примером фрактальной структуры. Фракталы — это круто и красиво, но эта статья не о них. Если вы хотите узнать больше о фракталах, прочтите прекрасную статью Спенсера Хукса. Я выбрал множество Мандельброта, потому что оно достаточно затратно в вычислительном отношении, но легко реализуемо и красиво! Так что же такое множество Мандельброта? Рассмотрим некоторую точку c на комплексной плоскости и простую итерацию

который отображает одну точку комплексной плоскости в другую. В качестве начального шага итерации выберите z₀=0. Поскольку n стремится к бесконечности, оказывается, что есть два варианта: либо точка уходит в бесконечность, либо остается близко к началу координат. Произойдет ли одно или другое, зависит от c. Если точка остается близкой к началу координат для всех n, то c является частью множества Мандельброта, в противном случае — нет. Это так просто. Давайте напишем простой код на Python для определения точек множества Мандельброта. Поскольку мы хотим оптимизировать производительность позже, нам нужно измерить время выполнения, поэтому я решил написать небольшой декоратор @measureдля удобства использования.

Я думаю, что код довольно понятен. Затем я могу запустить код и построить результаты с помощью matplotlib, вызвав

import matplotlib.pyplot as plt
X, Y, mandelbrot = compute_mandelbrot_set(1024, 768)
fig = plt.figure(figsize=(16, 8))
ax = fig.gca()
ax.set_aspect('equal')
ax.pcolor(X, Y, mandelbrot)

Это создает красивую фигуру из начала этой статьи и печатает необходимое время выполнения. На моем Макбуке написано:

compute_mandelbrot_set needed 12.472399 seconds.

Мы должны быть в состоянии заставить это работать быстрее! Если вы посмотрите на реализацию, энтузиасты Python сразу заметят, что может быть проблематичным: for-циклы, которые, как известно, в Python работают медленно. Вот жаль, потому что сам код вполне читабелен. Но прежде чем отказаться от читабельности, мы внесем несколько незначительных улучшений, которые могут немного помочь.

Первое, что я замечаю, это то, что мы дважды обращаемся к элементу массива Cs[irow, icol] в строках 31 и 33. Вместо этого нам следует определить скалярную переменную c = Cs[irow, icol] и использовать ее. Кажется, это не имеет большого значения, и на самом деле это не так, но если бы нам нужен доступ для чтения к другому большому массиву в том же месте кода, тогда это может иметь большое значение!

Причина в том, что когда вы читаете один элемент массива, процессор не извлекает из памяти только этот единственный элемент. Вместо этого он извлечет этот элемент и многие последующие элементы. Затем он сохранит их в кеше, предполагая, что последующие элементы могут понадобиться в ближайшее время. Если они нам понадобятся в ближайшее время, это будет огромным преимуществом в скорости, потому что процессору не нужно получать значения из ОЗУ, что очень-очень медленно (!), а вместо этого из кеша, что очень-очень быстро.

Но если мы читаем из какого-то массива a и сразу после этого из массива b, прежде чем вы снова получите доступ к a, то может случиться так, что элементы a в кэше были очищены и заменены элементами b. Итак, когда вам понадобится следующий элемент a, его уже нет в кеше! Это называется промахом кеша. В этом случае элемент приходится извлекать из оперативной памяти, а это, как мы уже говорили, медленно. Если он все еще в кеше, то у нас есть попадание в кеш, и это хорошо и быстро.

Векторизация

Что ж, как я уже сказал, циклы Python for печально известны своей медленностью. Вот почему у нас есть NumPy! Если мы работаем с массивами NumPy в целом, а не с отдельными элементами, то NumPy может использовать высокоэффективные параметры векторизации. Но это требует, чтобы мы полностью переписали наш код. К сожалению, во многих случаях это, мягко говоря, не делает код более читабельным. В нашем случае это будет выглядеть примерно так:

Запуск кода теперь возвращает:

compute_mandelbrot_set_2 needed 0.995727 seconds.

Неплохо, ускорение в 12 раз! Это сила векторизации NumPy. Но это пришлось расплачиваться за удобочитаемость. Чтобы мы могли работать со всеми массивами, мы должны определить неясные маски и еще более неясные логические операции над масками, чтобы получить тот же результат, что и с циклами for. Или кто-нибудь из вас считает это читаемым? Хорошо, это загадочно, и некоторые находят это крутым, но не могли бы вы сразу увидеть, что происходит? Я бы не стал. Так что, на мой вкус, цена за ускорение оказалась очень высокой. Возможно, у вас есть идея для более читаемого векторного решения. Если это так, пожалуйста, напишите комментарий! Но есть и другой вариант, сочетающий в себе лучшее из обоих миров: читабельность и скорость. Вот к чему мы пришли сейчас.

Лучшее из обоих миров

Вернемся к нашей исходной медленной реализации с for циклами. Самое смешное, что мы можем очень легко ускорить его, используя классный пакет Python под названием numba. Вы просто аннотируете свой код с помощью декораторов numba, и numba скомпилирует ваш декорированный код в чистый, оптимизированный машинный код. И оказывается, что результаты часто (не всегда) даже быстрее, чем векторизованный код NumPy. Единственное, что нам нужно сделать, это:

  • извлеките циклы for с тяжелыми вычислениями в отдельную функцию и
  • аннотировать новую функцию с помощью декоратора numba

Результат выглядит следующим образом:

Легко, не так ли? При первом запуске кода numba может потребоваться полсекунды для компиляции вашего кода, но при следующем запуске кода вы даже сэкономите эти полсекунды, если установите параметр cache=True в декораторе.

Итак, как это работает? На моем компе написано:

compute_mandelbrot_set_3 needed 0.581394 seconds.

Эй, в 2 раза быстрее, чем векторизованная версия NumPy. Неплохо, нумба!

Если вам понравилась эта статья, вы можете стать участником Medium, чтобы получить неограниченный доступ ко всем статьям Medium. Если вы зарегистрируетесь по этой ссылке, вы даже можете поддержать меня как автора, так как я получу долю в размере 2,27 доллара (на момент написания этой статьи) от платы за использование Medium.

Больше контента на plainenglish.io. Подпишитесь на нашу бесплатную еженедельную рассылку здесь.