numba - 성능 업!

파이썬을 사용하는 목적에 성능은 포함되어 있지 않다고 엄청 강조한다.
그렇지만 가끔은 성능이 좋아졌으면 하고 바랄 때가 있다.
결과를 보기 위해 몇 시간을 구동해야 한다면 말이다.

"High Performance Python(고성능 파이썬)"을 읽다 보니
책 제목답게 성능과 관련한 내용이 많았다.
당연히 나에게는 중요하지 않았다.
그런데, 파이썬 코드를 컴파일하는 챕터에서
거저 먹을 수 있는 몇 가지 방법을 알려주고 있었다.

numba 홈페이지에 가면 조금 무섭다.
괜히 설치되지 않을 것 같은 느낌!

  numba 홈페이지

정말 다행스럽게 파이참에서 바로 설치할 수 있었다.
어.. 일단 설치는 우분투 파이참에서 성공했다.
설정 들어가서 인터프리터로 이동한 다음 numba 검색 후에 설치.
혹시 맥이나 윈도우에서는 안될 수도 있겠다.
홈페이지의 분위기로 봐서는.
(안타깝게도 맥에서는 실패. 파이참에서 자동 설치가 되지 않는다.
터미널에서 설치할 때, llvmlite 모듈이 필요한데.. 이거 설치하면서 에러가 발생한다.
다른 사람들은 잘 되는 듯.. 나중에 다시 도전!!)

두 가지만 기억한다.

1. from numba import jit
   파일 상단에 numba를 import 하는데 사용하는 문법
2. @jit
   컴파일을 적용하려고 하는 함수 앞에 추가하는 문법
   numba는 jit(just in time) 컴파일러 기반이라서, 적용 함수의 실행 코드를 미리 생성한다는 뜻이다.

numba 홈페이지에 있는 코드 두 가지를 소개한다.
첫 번째는 기본 문법이 적용되는지 확인하기 위한 코드로
앞에서 언급한 jit를 import하고 @jit를 함수 앞에 추가한 부분만 보면 된다.

from numba import jit
from numpy import arange

# jit decorator tells Numba to compile this function.
# The argument types will be inferred by Numba when function is called.
@jit
def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i,j]
    return result

a = arange(9).reshape(3,3)
print(sum2d(a))

# [출력 결과]
# 36.0

두 번째 코드는 만델브로트 집합을 이미지로 표시하는 코드이다.
numba를 적용했을 때와 하지 않았을 때의 성능을 비교했다.
와우!!
26.5배 차이가 났다. 몇 줄 넣었을 뿐인데..
numba를 적용하기 위해 코드를 하나도 수정하지 않았는데..

#! /usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function, division, absolute_import

from timeit import default_timer as timer
from matplotlib.pylab import imshow, jet, show, ion
import numpy as np

from numba import jit


@jit
def mandel(x, y, max_iters):
    """
    Given the real and imaginary parts of a complex number,
    determine if it is a candidate for membership in the Mandelbrot
    set given a fixed number of iterations.
    """
    i = 0
    c = complex(x,y)
    z = 0.0j
    for i in range(max_iters):
        z = z*z + c
        if (z.real*z.real + z.imag*z.imag) >= 4:
            return i

    return 255

@jit
def create_fractal(min_x, max_x, min_y, max_y, image, iters):
    height = image.shape[0]
    width = image.shape[1]

    pixel_size_x = (max_x - min_x) / width
    pixel_size_y = (max_y - min_y) / height
    for x in range(width):
        real = min_x + x * pixel_size_x
        for y in range(height):
            imag = min_y + y * pixel_size_y
            color = mandel(real, imag, iters)
            image[y, x] = color

    return image

image = np.zeros((500 * 2, 750 * 2), dtype=np.uint8)
s = timer()
create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20)
e = timer()
print(e - s)
imshow(image)
# jet()
# ion()
show()

# [출력 결과]
# 적용 후 : 0.18618797000090126
# 적용 전 : 4.946132005999971

함수 앞에 적용한 @jit 데코레이터가 있을 때와 없을 때로 각각 구동했다.
적용했을 때는 결과가 바로 표시된다.
반면 적용하지 않았을 때는 꽤나 기다려야 된다.