numpy.einsum을 이용한 빠른 2d convolution 연산 (conv2d)예제 코드

 본 포스트는 numpy.einsum을 사용하여 2차원 합성곱 연산을 수행하는 예제 코드를 싣고 있다.

단순 연산을 사용한 2차원 합성곱

단순히 for loop를 사용하여 아래처럼 2차원 합성곱 연산을 수행할 수 있다.

import numpy as np

def conv2d(image, kernel):
    kernel = np.flipud(np.fliplr(kernel)) #XCorrel
    output = np.zeros(np.subtract(image.shape, kernel.shape))
    
    for y in range(output.shape[1]):
        for x in range(output.shape[0]):
            output[x, y] = (kernel * image[x: x + kernel.shape[0], y: y + kernel.shape[1]]).sum()

    return output

하지만, 위 경우 image의 크기가 커지면 연산양이 많아져 수행 속도가 느려지게 된다. 

numpy.einsum을 사용한 2차원 합성곱

 빠른 속도의 합성곱 연산을 원한다면, numpy.lib.stride_tricks.as_strided와 numpy.einsum을 사용하면 빠른 속도로 2차원 함성곱 연산을 수행할 수 있다.

import numpy as np

def conv2d_np(image, kernel):
    kernel = np.flipud(np.fliplr(kernel)) #XCorrel
    
    sub_matrices = np.lib.stride_tricks.as_strided(image,
                                                   shape = tuple(np.subtract(image.shape, kernel.shape))+kernel.shape, 
                                                   strides = image.strides * 2)

    return np.einsum('ij,klij->kl', kernel, sub_matrices)

numpy.lib.stride_tricks.as_strided는 배열을 주어진 shape과 stride에 맞게 새로운 배열로 만들어준다. 

https://numpy.org/doc/stable/reference/generated/numpy.lib.stride_tricks.as_strided.html

 

예를 들어 위 conv2d_np 코드에서는 입력 image가 아래와 같을 때,

ipdb> image
array([[11, 12, 13, 14, 15, 16],
       [21, 22, 23, 24, 25, 26],
       [31, 32, 33, 24, 25, 26],
       [41, 42, 43, 44, 45, 46],
       [51, 52, 53, 54, 55, 56]])

numpy.lib.stride_tricks.as_strided를 사용해 새로 만들어진 sub_matrices는 다음과 같다. 

ipdb> sub_matrices
array([[[[11, 12, 13],
         [21, 22, 23],
         [31, 32, 33]],

        [[12, 13, 14],
         [22, 23, 24],
         [32, 33, 24]],

        [[13, 14, 15],
         [23, 24, 25],
         [33, 24, 25]]],


       [[[21, 22, 23],
         [31, 32, 33],
         [41, 42, 43]],

        [[22, 23, 24],
         [32, 33, 24],
         [42, 43, 44]],

        [[23, 24, 25],
         [33, 24, 25],
         [43, 44, 45]]]])

즉 합성곱을 위해 image 배열을 kernel 배열의 크기에 맞게 분리해 새로운 배열로 만들어주는 일을 수행한다. 

위에서 만들어진 배열을 numpy.einsum을 사용하면 쉽게 스칼라곱을 취할 수 있다. 이는 다음의 코드를 numpy.einsum만으로 쉽고 빠르게 수행할 수 있다는 의미이다.

for y in range(output.shape[1]):
    for x in range(output.shape[0]):
        output[x, y] = (kernel * image[x: x + kernel.shape[0], y: y + kernel.shape[1]]).sum()

numpy.einsum은 아래 stackoverflow 사이트에 잘 설명되어 있다.

https://stackoverflow.com/questions/26089893/understanding-numpys-einsum

성능 비교

위 예제 코드 conv2d의 성능을 비교하면, 동일한 이미지에서 for loop를 사용한 conv2d는 약1.33초의 시간이 걸리고, numpy.einsum을 사용한 conv2d_np는 0.01초의 시간이 걸렸다.

numpy.einsum을 사용한 2차원 합성곱 연산이 약 133배 빠른 것을 볼 수 있다.