[numpy] np.sum()은 np.einsum()보다 느리다.

- np.einsum()이란?

numpy 라이브러리에는 각종 연산을 해주는 함수가 있습니다.

그중 아인슈타인 표기법을 함수로 만든 np.einsum()이 있습니다.

아인슈타인 표기법의 자세한 내용은 다음 링크를 참고해주세요. https://rockt.ai/2018/04/30/einsum

 

numpy 변수 관점에서 설명하면 아인슈타인 표기법은 axis(행과 열)에 따라 특정 규칙으로 계산을 하는 것입니다.

좀 더 간단하게 설명하기 위해 3x3 행렬을 생각해봅시다. 이는 3x3 shape의 numpy 배열이라 봐도 되겠죠.

행렬 안에 있는 원소 aij라치면 i는 행의 인덱스, j는 열의 인덱스입니다.

a31이면 행렬 내에서 3행 1열에 위치한다는 뜻이죠.

아인슈타인 표기법에서는 행렬 내의 원소를 ij와 같은 인덱스를 이용해서 행렬를 표기합니다.

 

만약 aij를 aji로 i와 j의 위치를 바꾼다고 치면 행과 열이 바뀌는 거죠? 이게 바로 np.transpose()입니다.

혹은 aij에 i=1부터 3까지, j=1부터 3까지 더하라고 하면 모든 원소의 합이 되고 이건 np.sum()입니다.

np.einsum()에 규칙을 어떻게 주느냐에 따라 np.transpose() 또는 np.sum()이 될수도 있습니다.

 

- np.einsum('ij->', vec2)와 np.sum(vec2)의 속도 차이

vec2는 n x n의 numpy 배열입니다. 제목의 두 함수는 완전히 같은 기능을 합니다.

하지만 이 둘의 실행시간은 다릅니다.

timeit로 실행시간을 확인해봅시다. 

# 100x100 numpy 배열 선언
dim = 100
vec2 = np.random.rand(dim, dim)

print("np.sum:", np.sum(vec2))
%timeit np.sum(vec2)
print("np.einsum:", np.einsum('ij->', vec2))
%timeit np.einsum('ij->', vec2)

"""
[Output]
np.sum: 5032.0468260231455
6.29 μs ± 79 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
np.einsum: 5032.046826023145
4.39 μs ± 43.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
"""

 

먼저 두 함수의 결과값을 확인하면 거의 같습니다(소수점 뒷자리 값이 좀 다릅니다).

np.sum()은 6.29 마이크로초, np.einsum()은 4.39 마이크로초가 걸렸습니다.

np.einsum()이 20~30% 정도 빠릅니다.

 

다음 그림은 배열의 크기를 바꿔가며 실행시간을 확인한 것입니다.

x축은 배열의 크기를, y축은 실행시간의 비율입니다.

y값이 작을수록 np.einsum()이 np.sum()보다 빠릅니다.

100x100 크기까지는 20% 정도 빠르다가 점점 배열의 크기가 커질수록 np.einsum()이 2배 이상 빨라집니다.

 

 

- 메모리 배열과 axis 기반 계산의 한계

왜 이런 현상이 벌어지는 지는 알아보기 위해 ChatGPT한테 물어보고, stackoverflow에도 검색을 해보았습니다.

여러 정보를 취합하고 제가 해석한 것이라 틀릴 수도 있습니다.

 

1. 메모리 배열

2차원 이상의 배열을 선언할 때 이 배열의 원소를 메모리에 저장하는 순서가 있습니다.

보통 두 종류의 순서가 있습니다. 행을 우선시 저장할 수도 있고, 열을 우선시 저장할 수도 있습니다.

행을 우선시 저장하는 것을 C-contiguous, 열을 우선시 저장하는 것을 F-Contiguous라고 부릅니다.

C는 C언어, F는 포트란(fortran) 언어를 의미합니다.

# 이렇게 보면 3x3 행렬
a = np.array([
   1열 2열 3열
1행 [1, 2 ,3], 
2행 [4, 5, 6], 
3행 [7, 8, 9]
])

a = np.array([[1, 2, 3] , [4, 5, 6], [7, 8, 9]]) 를 정의했다면 numpy는 기본 순서가 C-contiguous이므로

메모리 상에는 1 2 3 4 5 6 7 8 9 순서로 저장됩니다. 1행, 2행 3행 순서로 저장된거죠.

만약 F-contiguous로 저장한다면 1 4 7 2 5 8 3 6 9 순으로 1열, 2열, 3열 순서로 저장됩니다.

 

2. axis 기반 연산과 자료를 읽는 순서

np.sum()을 쓰면 axis 옵션으로 넣을 수 있습니다.

a = np.array([[1, 2, 3] , [4, 5, 6], [7, 8, 9]])
print(np.sum(a, axis=0))
# output: [12 15 18]

print(np.sum(a, axis=1))
#output: [ 6 15 24]

 

np.sum(a, axis=0)는 1열의 합, 2열의 합, 3열의 합,

np.sum(a, axis=1)은 1행의 합, 2행의 합, 3행의 합입니다.

 

np.sum()은 axis를 기반으로 동작하는데 axis에 따라서 계산시간도 다릅니다(배열의 크기가 작으면 거의 차이 없음).

import numpy as np
dim = 10000
vec2 = np.random.rand(dim, dim)

%timeit np.sum(vec2)
%timeit np.sum(vec2, axis=0)
%timeit np.sum(vec2, axis=1)

"""
NumPy는 기본적으로 C-contiguous
[output]
107 ms ± 1.95 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
47.3 ms ± 1.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
103 ms ± 2.91 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

 

axis=0일 때 47.3 마이크로 초, axis=1일 때는 103 마이크로 초로 axis=1일 때가 더 느립니다.

axis에 따라 실행시간 차이가 생기는 이유는 C-contiguous로 저장했기 때문에 메모리에 접근하는 방법이 axis에 따라 (열이나 행이냐에 따라)다르기 때문입니다.

전체평균인 np.sum(vec2)는 107 마이크로 초가 걸려 axis=1일때보다 조금 더 걸렸습니다.

 

numpy 라이브러리에는 배열을 F-contiguous로 저장하는 기능이 있습니다.

F-contiguous로 저장했을 땐 결과가 반대로 나와야겠죠?

import numpy as np
dim = 10000
vec2 = np.random.rand(dim, dim)

# fortran 형식의 array로 변환(F-contiguous)
vec2_f = np.asfortranarray(vec2)

%timeit np.sum(vec2_f)
%timeit np.sum(vec2_f, axis=0)
%timeit np.sum(vec2_f, axis=1)

"""
[output]
109 ms ± 1.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
106 ms ± 3.21 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
48.4 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

 

F-contiguous 배열에서는 np.sum(vec2, axis=0)일 때 시간이 더 오래걸립니다.

 

np.sum()은 axis 기반으로 동작하고 메모리에 접근할 때 특정 간격이 아니라 한 칸씩 자료를 불러들이면 시간이 오래걸린다고 합니다. 

3x3 행렬이 [1 2 3 4 5 6 7 8 9] 이런 순서(C-contiguous)로 메모리가 저장되어 있다고 치면 자료를

1, 4, 7 (3간격)부터 +1씩 2번 반복하는 것(np.sum(a, axis=0))이 1, 2, 3(한 칸씩) +3씩, 2번 반복하는 것 (np.sum(a, axis=1)) 보다 더 빠른듯합니다.

 

np.sum()은 axis 기반이라 느린 것은 맞고,

np.einsum()이 어떻게 연산을 하는 가에 대한 자세한 설명은 못 찾았는데 일단 axis 기반이 아니고 계산을 위한 더 효율적인 메모리 배치를 하여 최적화가 되어있다고 합니다.

 

- 마무리

이론상 np.einsum()이 무조건 빠를 것 같은데 꼭 그렇지는 않습니다.

제가 확인한 바로는 np.sum(), np.transpose(), np.diag() 기능을 np.einsum()으로 구현하면 더 빠릅니다.

하지만 저는 하나의 numpy 배열을 사용하고 단순한 계산을 하는데 np.einsum()을 사용했습니다.

여러 개의 배열을 사용하고 계산과정이 더 복잡한 경우에는 np.einsum()이 더 빠를 겁니다.