c++ - 使用 Intel MKL 的 3D 卷积

Question

我编写了一个 C/C++ 代码，它使用 Intel MKL 来计算一个数组的 3D 卷积，该数组的大约 300×200×200元素。我想应用一个 3×3×3 或 5×5×5 的内核。 3D 输入数组和内核都有实数值。

此 3D 数组以列方式存储为 double 类型的一维数组。类似地，内核的类型为 double 并按列保存。例如，

for( int k = 0; k < nk; k++ ) // Loop through the height.
    for( int j = 0; j < nj; j++ ) // Loop through the rows.
        for( int i = 0; i < ni; i++ ) // Loop through the columns.
        {
            ijk = i + ni * j + ni * nj * k;
            my3Darray[ ijk ] = 1.0;
        }

对于卷积的计算，我想对输入数组和内核执行 not-in-place FFT 并防止它们被修改(我需要稍后在我的代码中使用它们)然后进行反向计算 in-place。

当我比较从我的代码获得的结果与通过 MATLAB 获得的结果时，它们非常不同。有人可以帮我解决这个问题吗？我的代码中缺少什么？

这是我使用的 MATLAB 代码:

a = ones( 10, 10, 10 );
kernel = ones( 3, 3, 3 );
aconvolved = convn( a, kernel, 'same' );

这是我的 C/C++ 代码:

#include <stdio.h>
#include "mkl.h"

void Conv3D(
    double *in, double *ker, double *out,
    int nRows, int nCols, int nHeights)
{

    int NI = nRows;
    int NJ = nCols;
    int NK = nHeights;

    double *in_fft  = new double [NI*NJ*NK];
    double *ker_fft = new double [NI*NJ*NK];

    DFTI_DESCRIPTOR_HANDLE fft_desc = 0;
    MKL_LONG sizes[]   = { NK, NJ, NI };
    MKL_LONG strides[] = { 0, NJ*NI, NI, 1 };

    DftiCreateDescriptor( &fft_desc, DFTI_DOUBLE, DFTI_REAL, 3, sizes     );
    DftiSetValue        (  fft_desc, DFTI_PLACEMENT     , DFTI_NOT_INPLACE);   // Out-of-place computation.
    DftiSetValue        (  fft_desc, DFTI_INPUT_STRIDES , strides         );
    DftiSetValue        (  fft_desc, DFTI_OUTPUT_STRIDES, strides         );
    DftiSetValue        (  fft_desc, DFTI_BACKWARD_SCALE, 1/NI/NJ/NK      );
    DftiCommitDescriptor(  fft_desc );

    DftiComputeForward  (  fft_desc, in , in_fft  );
    DftiComputeForward  (  fft_desc, ker, ker_fft );

    for (long long i = 0; i < (long long)NI*NJ*NK; ++i )
        out[i] = in_fft[i]*ker_fft[i];

    // In-place computation.
    DftiSetValue        (  fft_desc, DFTI_PLACEMENT, DFTI_INPLACE );
    DftiCommitDescriptor(  fft_desc      );
    DftiComputeBackward (  fft_desc, out );

    DftiFreeDescriptor  ( &fft_desc );

    delete[] in_fft;
    delete[] ker_fft;

}

int main(int argc, char* argv[])
{
    int n = 10;
    int nkernel = 3;

    double *a          = new double [n*n*n]; // This array is real.
    double *aconvolved = new double [n*n*n]; // The convolved array is also real.
    double *kernel     = new double [nkernel*nkernel*nkernel]; // kernel is real.

    // Fill the array with some 'real' numbers.
    for( int i = 0; i < n*n*n; i++ )
        a[ i ] = 1.0;

    // Fill the kernel with some 'real' numbers.
    for( int i = 0; i < nkernel*nkernel*nkernel; i++ )
        kernel[ i ] = 1.0;

    // Calculate the convolution.
    Conv3D( a, kernel, aconvolved, n, n, n );

    printf("Convolved:\n");
    for( int i = 0; i < n*n*n; i++ )
        printf( "%15.8f\n", aconvolved[i] );

    delete[] a;
    delete[] kernel;
    delete[] aconvolved;

    return 0;
}

Answer 1

您无法使用实数值频率数据(仅幅度)反转 FFT。正向 FFT 需要输出复数数据。这是通过设置 DFTI_FORWARD_DOMAIN setting 来完成的。到 DFTI_COMPLEX。

DftiCreateDescriptor( &fft_desc, DFTI_DOUBLE, DFTI_COMPLEX, 3, sizes     );

这样做也隐式地将向后域设置为复数。

您还需要一个复杂的数据类型。大概是这样的，

MKL_Complex16* in_fft  = new MKL_Complex16[NI*NJ*NK];

这意味着您必须将实部和虚部相乘:

for (size_t i = 0; i < (size_t)NI*NJ*NK; ++i) {
    out_fft[i].real = in_fft[i].real * ker_fft[i].real;
    out_fft[i].imag = in_fft[i].imag * ker_fft[i].imag;
}

逆 FFT 的输出也很复杂，假设您的输入数据是真实的，您只需获取 .real 组件，这就是您的结果。这意味着您将需要一个临时的复杂输出数组(例如，如上所述的 out_fft)。

另请注意，为避免伪影，您希望每个维度上的 fft 大小(至少)为 M+N-1。通常，您会选择次高的 2 次幂来提高速度。

我强烈建议您首先使用 FFT 在 MATLAB 中实现它。有许多这样的实现可用 ( example )，但我会从基础开始，自己制作一个简单的函数。