Numpy 如何在一个轴上进行卷积

本文将介绍如何使用Numpy在一个轴上进行卷积。

在信号处理和图像处理中，卷积是一种基本运算，其作用是处理信号和图像。在Numpy中，我们可以使用numpy.convolve()函数来计算卷积。然而，numpy.convolve()函数同时计算多个轴的卷积，但在有些情况下，我们仅需要沿着某一轴上进行卷积。这时，我们可以使用Numpy的apply_along_axis()函数，该函数可以沿着一个指定轴应用一个函数。

阅读更多：Numpy 教程

单轴卷积

通常情况下，Numpy的convolve()函数会对输入数组的所有轴进行卷积，但这不总是我们所需要的。在某些情况下，我们仅需要沿着一个轴进行卷积，这时可以使用apply_along_axis()函数。

定义一个一维卷积核：

kernel = np.array([-1, 0, 1])

我们可以使用 apply_along_axis()函数对输入数组进行单轴卷积：

result = np.apply_along_axis(lambda m: np.convolve(m, kernel, mode='same'), axis=0, arr=input_array)

在上述代码中，lambda函数中的np.convolve()函数计算了沿着指定轴的卷积。axis=0参数指定了我们要沿着的轴。

考虑一组示例数据：

input_array = np.arange(12).reshape(3, 4)

这将创建一个3×4的数组：

array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

如果我们沿着第一轴对其进行卷积，可以计算出以下结果：

array([[  0,   1,   2,   1],
       [  4,   5,   6,   5],
       [  8,   9,  10,   9]])

这是通过将卷积核应用于每一行来得出的。

改进性能

当处理大型数组时，上述方法的运行速度可能很慢。为了提高性能，可以使用Numba库对我们的代码进行JIT编译。

先安装Numba库：

!pip install numba

这里是一个使用Numba加速的示例代码：

import numba

@numba.njit
def convolve_along_axis(x, kernel):
    """
    Compute 1D convolution over specified axis of array
    """
    n_x, n_k = x.shape[0], kernel.shape[0]
    pad_x = np.pad(x, ((0, n_k-1), (0, 0)), mode='constant')
    out = np.zeros(pad_x.shape)
    for i in range(n_x):
        out[i] = np.sum(pad_x[i:i+n_k] * kernel, axis=0)
    return out[:n_x]

def run_convolution(input_array, kernel, axis=0):
    """
    Apply kernel along specified axis of input_array
    """
    return np.apply_along_axis(convolve_along_axis, axis=axis, arr=input_array, kernel=kernel)

这里，我们编写了一个convolve_along_axis()函数来处理卷积。使用np.pad()函数，我们将输入数组界定在合适大小的帧内。正如前面所提到的，Numba库可以加速使用JIT编译的函数。由于我们的函数采用了Numba装饰器，因此Numba将启用JIT编译。这样，我们就可以通过run_convolution()函数调用apply_along_axis()函数来计算沿着指定轴的卷积。