如何使用Python中的GPU执行一维FFT?

时间:2019-05-31 17:44:14

标签: python gpu fft pycuda

嗨,我正在尝试使用python中的GPU做一个简单的一维fft。我已经找到了2D的示例代码:

import numpy as np
import pycuda.autoinit
import pycuda.gpuarray as gpuarray
import skcuda.fft as cu_fft

 def fft2_gpu(x, fftshift=False):

    ''' This function produce an output that is 
    compatible with numpy.fft.fft2
    The input x is a 2D numpy array'''

    # Convert the input array to single precision float
    if x.dtype != 'float32':
        x = x.astype('float32')

    # Get the shape of the initial numpy array
    n1, n2 = x.shape

    # From numpy array to GPUarray
    xgpu = gpuarray.to_gpu(x)

    # Initialise output GPUarray 
    # For real to complex transformations, the fft function computes 
    # N/2+1 non-redundant coefficients of a length-N input signal.
    y = gpuarray.empty((n1,n2//2 + 1), np.complex64)

    # Forward FFT
    plan_forward = cu_fft.Plan((n1, n2), np.float32, np.complex64)
    cu_fft.fft(xgpu, y, plan_forward)

    left = y.get()

    # To make the output array compatible with the numpy output
    # we need to stack horizontally the y.get() array and its flipped version
    # We must take care of handling even or odd sized array to get the correct 
    # size of the final array   
    if n2//2 == n2/2:
        right = np.roll(np.fliplr(np.flipud(y.get()))[:,1:-1],1,axis=0)
    else:
        right = np.roll(np.fliplr(np.flipud(y.get()))[:,:-1],1,axis=0) 

    # Get a numpy array back compatible with np.fft
    if fftshift is False:
        yout = np.hstack((left,right))
    else:
        yout = np.fft.fftshift(np.hstack((left,right)))

    return yout.astype('complex128')

我试图像这样对1D进行修改。我很迷茫,不熟悉ffts:

def fft_gpu(x, fftshift=False):

    ''' This function produce an output that is 
    compatible with numpy.fft.fft2
    The input x is a 2D numpy array'''

    # Convert the input array to single precision float
    if x.dtype != 'float32':
        x = x.astype('float32')

    # Get the shape of the initial numpy array
    n1 = len(x)

    # From numpy array to GPUarray
    xgpu = gpuarray.to_gpu(x)

    # Initialise output GPUarray 
    # For real to complex transformations, the fft function computes 
    # N/2+1 non-redundant coefficients of a length-N input signal.
    y = gpuarray.empty((n1), np.complex64)

    # Forward FFT
    plan_forward = cu_fft.Plan((n1), np.float32, np.complex64)
    cu_fft.fft(xgpu, y, plan_forward)

    left = y.get()
    """
    # To make the output array compatible with the numpy output
    # we need to stack horizontally the y.get() array and its flipped version
    # We must take care of handling even or odd sized array to get the correct 
    # size of the final array   
    if n2//2 == n2/2:
        right = np.roll(np.fliplr(np.flipud(y.get()))[:,1:-1],1,axis=0)
    else:
        right = np.roll(np.fliplr(np.flipud(y.get()))[:,:-1],1,axis=0) 
    """



    # Get a numpy array back compatible with np.fft
    if fftshift is False:
        yout = np.asarray(left)
    else:
        yout = np.fft.fftshift(np.hstack((left,left)))

    return yout.astype('complex128')

它给我的结果与我使用numpy fft时相似,但是到数组末尾的结果却有所不同,其结果是一堆零条目。对于我的1D fft函数中要更改的内容的任何见解,将不胜感激。

0 个答案:

没有答案