성능 및 코드 간결성 비교 numba.vectorize vs numba.jit
Post

성능 및 코드 간결성 비교 numba.vectorize vs numba.jit

간편하면서도 강력한 Python JIT 도구 numba의 맨 마지막에 언급한 바와 같이 numba.vectorize를 이용하여 생성한 ufunc는 element-by-element fashion 때문에 몇 가지 제약이 존재한다. 그 중 여기서 언급할 것 중 하나는 출력의 dimension이 입력의 dimension과 동일하다는 제약이다. 흑백 이미지 한 장입 입력받고 칼라 이미지를 출력하는 기능을 수행하는 함수를 numba 를 이용하여 구현하고자 할 때, vectorize를 이용하여 생성된 ufunc의 단일 호출로는 불가능하다. 이때는 numba.jit 를 사용해야 한다.

아래에 vectorize 를 이용한 func1과 jit를 이용한 func2를 제시한다. func1은 매개변수로 들어온 x, y에 대해 atan2 와 hypot를 계산한 후 두 값을 합한다. 특별한 의미가 있는 연산이 아니다. 단지 func2는 func1과는 다르게 출력의 dimension이 입력의 dimension과 같지 않을 수 있음을 보여주고자 하는 것이다. 예를 들어 설명하면 func2는 1채널 이미지 2장을 입력 받아 2채널 이미지 1장을 출력하는 함수이다. 첫 번째 채널은 atan2를 계산한 결과이고 두 번째 채널은 atan2 + hypot의 결과이다. 연산이 atan2, hypot, sum 총 3가지 이므로 속도 비교를 위해 func1에서도 이와 동일한 연산을 수행하게 하였다(func2 는 2채널에 대한 메모리 접근을 수행해야 하므로 func1 보다는 약간의 오버헤드가 있을 수 있음).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import math

import numpy
from numba import jit, vectorize, float32, prange


@vectorize([float32(float32, float32)], target='parallel')
def func1(x, y):
    a = math.atan2(x, y)
    b = math.hypot(x, y)
    return a + b

@jit(nopython=True, parallel=True)
def func2(x, y):
    n = x.shape[0]
    res = numpy.empty((2, n), dtype=x.dtype)
    for i in prange(n):
        a = math.atan2(x[i], y[i])
        b = math.hypot(x[i], y[i])
        res[0, i] = a
        res[1, i] = a + b
    return res

아래는 1000x1000부터 10000x10000까지 배열의 element 개수를 변경하면서 수행 속도를 측정한 결과이다. vectorize가 jit 보다 근소하게 빠르게 수행됨을 볼 수 있다. 실제 성능의 차이인지 2채널 접근으로 인한 차이인지 확인해볼 필요가 있겠지만 이에 대해서는 ‘2채널 접근에 약간의 비용이 소요될 것이다’ 정도로 추측하고 넘어가겠다. 특이할 점은 jit의 경우 최초 실행 시 수행 시간이 굉장히 길 다는 것이다(아래 그림의 맨 좌측 청색 점). JIT compilation 수행에 소요되는 시간을 것으로 추측되는데 함수 호출 시점이 아닌 함수 로드 시점에 compilation을 수행하는 방법을 찾아볼 필요가 있겠다.

img-01

이 정도의 성능과 간결성이면 시간을 들여서라도 기존의 코드들을 numba 방식으로 전환할 이유가 있어 보인다. 몇 가지 특이적 경우에 대한 샘플을 더 고려해 본 후 결정해야 겠다.

간편하면서도 강력한 Python JIT 도구 numba

numba를 이용한 Single-CPU, Multi-CPU, GPU-CUDA의 box blur 속도 비교