Numba 5분 가이드 — @jit의 두 모드와 성능 측정

Numba 5분 가이드

Numba는 NumPy 배열·함수·루프를 많이 쓰는 코드에 가장 잘 맞아요. 가장 흔한 사용법은 데코레이터를 함수에 붙여 컴파일을 지시하는 거예요. 호출 시점에 "just-in-time"으로 머신코드로 컴파일되고, 이후에는 전부 또는 일부가 네이티브 머신코드 속도로 돌아요.

출처: https://numba.readthedocs.io/en/stable/user/5minguide.html

설치

conda install numba
pip install numba

object 모드 vs nopython 모드

@jit 데코레이터는 기본적으로 nopython 모드로 동작해요. 이 모드는 데코레이터된 함수를 파이썬 인터프리터 개입 없이 온전히 실행되도록 컴파일하므로 최고 성능을 내는 권장 방식이에요. nopython 모드 컴파일이 실패하면 object 모드로 컴파일할 수 있는데, forceobj=True로 지정해요.

from numba import jit
import pandas as pd

x = {'a': [1, 2, 3], 'b': [20, 30, 40]}

@jit(forceobj=True, looplift=True)  # object mode 사용, 루프만 컴파일 시도
def use_pandas(a):
    # 파이썬 객체 위주
    pass

looplift=True는 루프를 머신코드 함수로 컴파일하려 시도해 순수 object 모드보다 나을 수 있어요.

성능 측정 주의

Numba는 주어진 인자 타입에 대해 함수를 먼저 컴파일해야 해요. 컴파일은 시간이 걸리지만, 한 번 컴파일되면 캐시해서 같은 타입으로 다시 부르면 재사용해요. 첫 호출에 컴파일 시간이 포함돼 측정을 왜곡할 수 있으니 주의해야 해요.

from numba import jit
import numpy as np
import time

x = np.arange(100).reshape(10, 10)

@jit(nopython=True)
def go_fast(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

start = time.perf_counter()
go_fast(x)          # 컴파일 포함
end = time.perf_counter()
print("Elapsed (with compilation) = {} s".format((end - start)))

start = time.perf_counter()
go_fast(x)          # 캐시에서 실행
end = time.perf_counter()
print("Elapsed (after compilation) = {} s".format((end - start)))

얼마나 빠른가

nopython 모드(또는 루프 컴파일)가 되면 특정 CPU에 맞춰 컴파일해요. 애플리케이션에 따라 다르지만 보통 1~2자릿수(order of magnitude) 속도 향상을 볼 수 있어요.

원리

Numba는 데코레이터된 함수의 파이썬 바이트코드를 읽고 입력 인자 타입 정보와 결합해 분석·최적화 후, LLVM 컴파일러 라이브러리로 CPU에 맞춘 머신코드를 생성해요. 컴파일된 버전은 이후 매 호출에 사용돼요.

다른 데코레이터

@njit — @jit(nopython=True)의 별칭(매우 흔함)
@vectorize — NumPy ufunc 생성
@guvectorize — 일반화 ufunc 생성
@stencil — 스텐실 연산 커널
@jitclass — JIT 인지 클래스
@cfunc — C/C++에서 호출할 네이티브 콜백
@overload — nopython 모드용 함수 구현 등록

Numba 5분 가이드 — @jit의 두 모드와 성능 측정