[Python] 연산 속도 올리기 - Numba

🐍 본 게시글은 Python 3.11.0 환경에서 작성되었습니다!

 

Numba는 Python을 위한 Just-In-Time(JIT) 컴파일러로, 주로 수치 계산 및 과학적 컴퓨팅을 최적화하는 데 사용된다.

Numpy와 함께 사용될 때 특히 강력하다.

 

특징

• JIT 컴파일 : 파이썬 코드를 실행 시간에 컴파일 하여 성능을 향상한다.
• Numpy 호환 : Numpy 배열과 많은 함수와 연산을 지원한다.
• GPU 지원 : NVIDIA GPU에서 CUDA를 사용하여 병렬 처리를 지원한다.
• 간편한 사용 : 기존 Python 코드에 주석을 추가하는 방식으로 쉽게 사용할 수 있다.

 

장점

• 성능 향상 : Python 코드를 C, Fortran과 유사한 속도로 실행 할 수 있게 한다.
• 직관적인 인터페이스 : 코드를 최소한의 수정으로 사용할 수 있어, 학습 곡선이 낮다.
• 다양한 기능 : JIT 컴파일 외에도, 병렬 처리, GPU 가속, 다양한 데이터 타입 지원 등 다양한 기능을 제공한다.

 

GitHub - numba/numba: NumPy aware dynamic Python compiler using LLVM

 

pip install numba

 

import numpy as np
from numba import njit
import time

def sum_array1(arr):
    total = 0.0
    for i in range(arr.size):
        total += arr[i]
    return total

@njit
def sum_array2(arr):
    total = 0.0
    for i in range(arr.size):
        total += arr[i]
    return total

# NumPy 배열 생성
data = np.arange(1, 100000001)

# Numba를 사용한 함수 호출
start = time.time()
result = sum_array1(data)
print(f"No JIT:{time.time() - start:.3f}sec")
start = time.time()
result = sum_array2(data)
print(f"JIT:{time.time() - start:.3f}sec")

대용량 연산의 경우 JIT를 사용하면 확실하게 빠르지만, 오히려 작은 연산의 경우 초기 컴파일 시간이 있으므로 호히려 성능이 떨어질 수 있다.

 

그리고 Numba는 모든 Python lib를 지원하지 않는다. (OpenCV 지원하지 않음...)

 

데코레이터

@njit

• Numba의 가장 기본적인 컴파일러

from numba import njit

@njit
def my_function(x):
    return x + 1

 

@jit

• @njit와 동일하나, 컴파일 옵션을 추가 할 수 있는 인자를 가진다. no_python=True를 추가하는 경우 Python 객체를 사용하지 않도록 강제한다.

from numba import jit

@jit(nopython=True)
def my_function(x):
    return x + 1

 

@cuda.jit

• CUDA GPU에서 실행 할 수 있는 커널 함수를 정의하는 데 사용된다. GPU 병렬 처리를 가능하게 한다.

from numba import cuda

@cuda.jit
def gpu_function(a, b, c):
    idx = cuda.grid(1)
    if idx < c.size:
        c[idx] = a[idx] + b[idx

 

@guvectorize

• 벡터화된 함수를 정의할 때 사용, 입력 배열의 형태를 지정하고, 다양한 크기의 배열에 대해 동일한 함수가 작동하도록 한다.

from numba import guvectorize

@guvectorize(["(float64[:], float64[:])"], "(n)->(n)")
def vectorized_function(input_array, output_array):
    for i in range(input_array.shape[0]):
        output_array[i] = input_array[i] * 2

 

@vectorize

• 배열 연산을 벡터화하여 실행하는 데 사용한다. 스칼라 함수를 벡터화하여 입력 배열에 대해 실행 할 수 있다.

from numba import vectorize

@vectorize([float64(float64)])
def vectorized_scalar_function(x):
    return x * 2

 

@stencil

• 격자 데이터에 대해 스텐실 연산을 수행할 때 사용된다. 주로 이미지 처리나 과학적 계산에서 사용된다.

from numba import stencil

@stencil
def stencil_function(a):
    return a[0, 0] + a[1, 0] + a[0, 1] + a[1, 1]

 

@prange

• 병렬 처리를 위한 반복문을 작성할 때 사용된다. for루프를 병렬로 실행 할 수 있게 해준다.

from numba import njit, prange

@njit
def parallel_sum(arr):
    total = 0.0
    for i in prange(arr.size):
        total += arr[i]
    return total