Python에서 2차원 합성곱 빠르게 계산하기
영상 처리에서 2차원 합성곱(2-D convolution)은 매우 유용한 연산이다. Blurring, morphology, edge 검출, sharpening 등과 같은 작업에 널리 활용된다. Python의 단순한 2차원 합성곱 방법은 큰 이미지에서 막대한 계산 부하를 가진다. 이 글에서는 합성곱 알고리즘의 성능을 높이기 위해 stride trick을 사용하는 방법을 소개한다. 이 방법은 메모리에 직접 접근하므로, 메모리에 데이터를 쓸 때는 주의해야 한다.
2차원 합성곱의 기본적인 구현은 다음 코드로 나타낼 수 있다.
import numpy as np
import time
def conv2d(x, kernel):
H, W = x.shape
H_k, W_k = kernel.shape
kernel_width_half = H_k // 2
out = np.zeros_like(x)
x_pad = np.pad(x, kernel_width_half)
for i in range(0, H):
for j in range(0, W):
for m in range(0, H_k):
for n in range(0, W_k):
out[i][j] += x_pad[i+m][j+n] * kernel[m][n]
return out
이 방법은 여러 겹의 중첩 for 루프를 포함하므로, 입력 행렬의 크기가 커질수록 계산 시간이 증가한다. 예를 들어 320x240 이미지와 3x3 커널의 경우, 이 2차원 합성곱 알고리즘은 M1 프로세서가 탑재된 MacBook Air에서 약 500 ms가 걸린다.
2차원 합성곱 알고리즘의 성능을 개선하기 위해 np.lib.stride_tricks.as_strided()를 활용할 수 있다. 본질적으로 이 함수는 원하는 shape과 stride를 가진 새로운 배열을 만든다. stride란 특정 차원에서 한 단계 이동할 때 건너뛰는 바이트 수를 가리킨다. 예를 들어 numpy 배열 x = np.array([[1,2],[3,4]], dtype=np.int64)를 생각해 보자. x의 원시 데이터 바이트는 다음과 같이 볼 수 있다.
x = np.array([[1,2],[3,4]], dtype=np.int64)
print(x.tobytes()) # little-endian
# b'\x01\x00\x00\x00\x00\x00\x00\x00 \x02\x00\x00\x00\x00\x00\x00\x00 \x03\x00\x00\x00\x00\x00\x00\x00 \x04\x00\x00\x00\x00\x00\x00\x00'
여기서 x의 열 축 방향 stride가 8바이트이기 때문에, x[0]은 b'\x01\x00\x00\x00 \x00\x00\x00\x00'에, x[1]은 b'\x02\x00\x00\x00 \x00\x00\x00\x00'에 대응한다. 같은 방식으로 행 축 방향 stride는 16바이트다.
print(x.strides)
# (16, 8)
따라서 np.lib.stride_tricks.as_strided(x, (4,), (8,))는 [1,2,3,4]를 반환하고, np.lib.stride_tricks.as_strided(x, (4,), (4,))는 [1, 8589934592, 2, 12884901888]을 반환한다. 여기서 b'\x00\x00\x00\x00 \x02\x00\x00\x00'은 \(256^4 \times 2 = 8,589,934,592\)에, b'\x00\x00\x00\x00 \x03\x00\x00\x00'은 \(256^4 \times 3 = 12,884,901,888\)에 대응한다.
다시 2차원 합성곱 방법으로 돌아가서, \(H \times W\) 행렬 x와 \(M \times N\) 커널 kernel에 대해, x_pad = np.pad(x, ((M//2, M//2), (N//2, N//2)))로 행렬에 패딩을 한 뒤 np.lib.stride_tricks.as_strided(x_pad, (H, W, M, N), x_pad.strides + x_pad.strides)를 사용하면 \(H \times W \times M \times N\) 크기의 텐서를 만든다. 세 번째와 네 번째 차원으로 이루어진 각 2차원 부분 행렬은 커널과 원소별로 곱해진다. 최종 출력 행렬은 아인슈타인 합 규칙(Einstein summation convention)을 사용하여 np.einsum('klij,ij->kl', sub_matrices, kernel)로 얻는다. 첫 번째 인자에서 klij는 sub_matrices 텐서의 1~4번째 차원을, ij는 커널 행렬의 차원을 나타낸다. 따라서 klij,ij->kl은 다음을 의미한다.
여기서 \(S\)는 sub_matrices 텐서, \(K\)는 kernel 행렬, \(O\)는 출력 행렬이다. 개선된 2차원 합성곱 전체 코드는 다음과 같다.
def fast_conv2d(x, kernel):
H, W = x.shape
H_k, W_k = kernel.shape
kernel_width_half = H_k // 2
x_pad = np.pad(x, kernel_width_half)
sub_matrices = np.lib.stride_tricks.as_strided(x_pad, (H, W, H_k, W_k), x_pad.strides*2)
return np.einsum('klij,ij->kl', sub_matrices, kernel)
또한, 위 코드에서 np.lib.stride_tricks.as_strided(x_pad, (H, W, H_k, W_k), x_pad.strides*2)는 np.lib.stride_tricks.sliding_window_view(x_pad, (H_k, W_k))로 대체할 수 있다.
def fast_conv2d(x, kernel):
H_k, W_k = kernel.shape
kernel_width_half = H_k // 2
x_pad = np.pad(x, kernel_width_half)
sub_matrices = np.lib.stride_tricks.sliding_window_view(x_pad, (H_k, W_k))
return np.einsum('klij,ij->kl', sub_matrices, kernel)
위 세 가지 방법의 계산 시간을 비교하면, 뒤의 두 방법이 훨씬 빠름을 알 수 있다. 또한 as_strided()가 sliding_window_view()보다 약간 더 좋은 성능을 보인다. 아래 그래프의 오차 막대는 표준편차의 3배를 의미한다.
