[英]GPU FFT Convolution using Cupy
我正在嘗試使用 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。
有人可以指出我正確的方向嗎?
解決方案是使用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
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.