GPU FFT 卷积使用 Cupy

Question

我正在尝试使用Cupy在GPU上执行FFT卷积操作。

使用scipy.signal.fftconvolve的源代码，我想出了以下基于Numpy的函数，它运行良好：

import numpy as np
def FFTConvolve(in1, in2):
if in1.ndim == in2.ndim == 0:  # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("Dimensions do not match.")
elif in1.size == 0 or in2.size == 0:  # empty arrays
return array([])
s1 = np.asarray(in1.shape)
s2 = np.asarray(in2.shape)
shape = s1 + s2 - 1
fsize = 2 ** np.ceil(np.log2(shape)).astype(int) 
fslice = tuple([slice(0, int(sz)) for sz in shape])
ret = np.fft.ifft(np.fft.fft(in1, fsize) * np.fft.fft(in2, fsize))[fslice].copy()
return ret

我天真地为Cupy编写了如下程序：

import cupy as cp
def FFTConvolve(in1, in2):
if in1.ndim == in2.ndim == 0:  # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("Dimensions do not match.")
elif in1.size == 0 or in2.size == 0:  # empty arrays
return array([])
in1 = cp.asarray(in1)
in2 = cp.asarray(in2)
s1 = cp.asarray(in1.shape)
s2 = cp.asarray(in2.shape)
shape = s1 + s2 - 1
fsize = 2 ** cp.ceil(cp.log2(shape)).astype(int) 
fslice = tuple([slice(0, int(sz)) for sz in shape])
ret = cp.fft.ifftn(cp.fft.fftn(in1, fsize) * cp.fft.fftn(in2, fsize))[fslice].copy()
return ret

后者在第enter code here行上给了我以下错误：

TypeError: 'cupy.core.core.ndarray' object cannot be interpreted as an integer

cupy.fft.ftt 的文档指出它接受元组作为一个范围，但由于某种原因将其读取为 cupy.ndarray。

有人可以指出我正确的方向吗？

Answer 1

解决方案是使用cp.asnumpy()命令：

def FFTConvolve(in1, in2):
if in1.ndim == in2.ndim == 0:  # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("Dimensions do not match.")
elif in1.size == 0 or in2.size == 0:  # empty arrays
return array([])
s1 = np.asarray(in1.shape)
s2 = np.asarray(in2.shape)
shape = s1 + s2 - 1
fsize = 2 ** np.ceil(np.log2(shape)).astype(int) 
fslice = tuple([slice(0, int(sz)) for sz in shape])
ret = cp.fft.ifftn(cp.fft.fftn(in1, np.asarray(fsize)) * cp.fft.fftn(in2, np.asarray(fsize)))[fslice].copy()
return ret