scipy 仅沿着一个轴卷积

yquaqz18  于 2022-12-04  发布在  其他
关注(0)|答案(5)|浏览(159)

我有两个第一轴维度相同的2-D数组。在python中,我只想沿着第二轴卷积这两个矩阵。我想得到下面的C,而不计算沿着第一轴的卷积。

import numpy as np
import scipy.signal as sg

M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

C = sg.convolve(A, B, 'full')[(2*M-1)/2]

有捷径吗?

x7yiwoj4

x7yiwoj41#

您可以使用np.apply_along_axis沿着所需的轴应用np.convolve。下面是一个将矩形波串滤镜应用于二维数组的示例:

import numpy as np

a = np.arange(10)
a = np.vstack((a,a)).T

filt = np.ones(3)

np.apply_along_axis(lambda m: np.convolve(m, filt, mode='full'), axis=0, arr=a)

这是一种简单的方法,可以推广许多没有axis参数的函数。

3yhwsihp

3yhwsihp3#

np.apply_along_axis实际上对你没有什么帮助,因为你试图迭代 two 数组,实际上,你必须使用一个循环,如here所述。
现在,如果数组很小,循环就可以了,但如果N和P很大,那么您可能需要使用FFT来进行卷积。
但是,您需要首先对数组进行适当的零填充,以便“完整”卷积具有预期的形状:

M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

A_ = np.zeros((M, N+P-1), dtype=A.dtype)
A_[:, :N] = A
B_ = np.zeros((M, N+P-1), dtype=B.dtype)
B_[:, :P] = B

A_fft = np.fft.fft(A_, axis=1)
B_fft = np.fft.fft(B_, axis=1)
C_fft = A_fft * B_fft

C = np.real(np.fft.ifft(C_fft))

# Test
C_test = np.zeros((M, N+P-1))
for i in range(M):
    C_test[i, :] = np.convolve(A[i, :], B[i, :], 'full')

assert np.allclose(C, C_test)
voj3qocg

voj3qocg4#

对于2D数组,函数scipy.signal.convolve2d更快,scipy.signal.fftconvolve甚至更快(取决于数组的维数):
这里同样代码用N = 100000

import time    
import numpy as np
import scipy.signal as sg

M, N, P = 10, 100000, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

T1 = time.time()
C = sg.convolve(A, B, 'full')
print(time.time()-T1)

T1 = time.time()
C_2d = sg.convolve2d(A, B, 'full')
print(time.time()-T1)

T1 = time.time()
C_fft = sg.fftconvolve(A, B, 'full')
print(time.time()-T1)

>>> 12.3
>>> 2.1
>>> 0.6

答案都是一样的,只是由于使用了不同的计算方法而略有不同(例如,fft与直接乘法,但我不知道convolve2d使用的是什么):

print(np.max(np.abs(C - C_2d)))
>>>7.81597009336e-14

print(np.max(np.abs(C - C_fft)))
>>>1.84741111298e-13
ltqd579y

ltqd579y5#

迟来的答案,但值得张贴参考。引用OP的评论:
A中的每一行都被B中对应的行过滤。我可以这样实现它,只是觉得可能有更快的方法。
A的大小大约是10千兆字节,我使用重叠相加。

天真/直接的方法

import numpy as np
import scipy.signal as sg

M, N, P = 4, 10, 20
A = np.random.randn(M, N) # (4, 10)
B = np.random.randn(M, P) # (4, 20)

C = np.vstack([sg.convolve(a, b, 'full') for a, b in zip(A, B)])

>>> C.shape
(4, 29)

A中的每一行与B中的每一相应行卷积,实质上卷积M个1D阵列/向量。

无循环+ CUDA支持的版本

可以使用PyTorch的F.conv1d来复制此操作。我们必须将A想象为一个长度为10* 的 *4通道一维信号。我们希望将A中的每个通道与长度为20的特定内核进行卷积。这是一种称为深度卷积的特殊情况,通常用于深度学习。
请注意,torch的卷积是以互相关的形式实现的,因此我们需要提前翻转B以进行实际卷积。

import torch
import torch.nn.functional as F

@torch.no_grad()
def torch_conv(A, B):
    M, N, P = A.shape[0], A.shape[1], B.shape[1]
    C = F.conv1d(A, B[:, None, :], bias=None, stride=1, groups=M, padding=N+(P-1)//2)
    return C.numpy()

# Convert A and B to torch tensors + flip B

X = torch.from_numpy(A) # (4, 10)
W = torch.from_numpy(np.fliplr(B).copy()) # (4, 20)

# Do grouped conv and get np array
Y = torch_conv(X, W)

>>> Y.shape
(4, 29)

>>> np.allclose(C, Y)
True

使用焊炬深度回旋的优点:

  • 没有循环!
  • 上面的解决方案也可以在CUDA/GPU上运行,如果AB是非常大的矩阵,那么CUDA/GPU确实可以加快速度。(从OP的评论来看,情况似乎是这样的:A的大小为10 GB。)

缺点:

  • 从数组转换为Tensor的开销(应忽略不计)
  • 操作前需要翻转B一次

相关问题