성능 및 코드 간결성 비교 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 보다는 약간의 오버헤드가 있을 수 있다)

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 을 수행하는지 모르겠다.(아니면 최초의 실행에 소요되는 시간이 현재로서는 알 수 없는 다른 무언가로 인한 것을 수도 있다)

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