GPU FFT Convolution с использованием Cupy - PullRequest
Новая
       100

GPU FFT Convolution с использованием Cupy

0 голосов
/ 06 января 2020

Я пытаюсь использовать Cupy для выполнения операции свертки 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.

Может кто-нибудь любезно указать мне в правильном направлении?

1 Ответ

0 голосов
/ 06 января 2020

Решением было использование команды 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
...