torch.rfft - 基于 fft 的卷积创建与空间卷积不同的输出

Question

我在 Pytorch 中实现了基于 FFT 的卷积，并通过 conv2d() 函数将结果与空间卷积进行了比较。使用的卷积滤波器是平均滤波器。 conv2d() 函数由于按预期进行平均滤波而产生了平滑的输出，但基于 fft 的卷积返回了更加模糊的输出。我在这里附上了代码和输出 -

空间卷积-

from PIL import Image, ImageOps
import torch
from matplotlib import pyplot as plt
from torchvision.transforms import ToTensor
import torch.nn.functional as F
import numpy as np

im = Image.open("/kaggle/input/tiger.jpg")
im = im.resize((256,256))
gray_im = im.convert('L') 
gray_im = ToTensor()(gray_im)
gray_im = gray_im.squeeze()

fil = torch.tensor([[1/9,1/9,1/9],[1/9,1/9,1/9],[1/9,1/9,1/9]])

conv_gray_im = gray_im.unsqueeze(0).unsqueeze(0)
conv_fil = fil.unsqueeze(0).unsqueeze(0)

conv_op = F.conv2d(conv_gray_im,conv_fil)

conv_op = conv_op.squeeze()

plt.figure()
plt.imshow(conv_op, cmap='gray')

基于 FFT 的卷积 -

def fftshift(image):
    sh = image.shape
    x = np.arange(0, sh[2], 1)
    y = np.arange(0, sh[3], 1)
    xm, ym  = np.meshgrid(x,y)
    shifter = (-1)**(xm + ym)
    shifter = torch.from_numpy(shifter)
    return image*shifter

shift_im = fftshift(conv_gray_im)
padded_fil = F.pad(conv_fil, (0, gray_im.shape[0]-fil.shape[0], 0, gray_im.shape[1]-fil.shape[1]))
shift_fil = fftshift(padded_fil)
fft_shift_im = torch.rfft(shift_im, 2, onesided=False)
fft_shift_fil = torch.rfft(shift_fil, 2, onesided=False)
shift_prod = fft_shift_im*fft_shift_fil
shift_fft_conv = fftshift(torch.irfft(shift_prod, 2, onesided=False))

fft_op = shift_fft_conv.squeeze()
plt.figure('shifted fft')
plt.imshow(fft_op, cmap='gray')

原图-

空间卷积输出-

基于fft的卷积输出-

有人可以解释一下这个问题吗？

Answer 1

您的代码的主要问题是 Torch 不计算复数，其 FFT 的输出是一个 3D 数组，第 3 维有两个值，一个用于实部，一个用于虚部。因此，乘法不进行复数乘法。

目前 Torch 中没有定义复数乘法（参见 this issue），我们必须自己定义。

一个小问题，但如果你想比较两个卷积操作也很重要，如下：

FFT 在第一个元素（图像的左上角像素）中获取其输入的原点。为避免移位输出，您需要生成一个填充内核，其中内核的原点是左上角的像素。这很棘手，实际上...

您当前的代码：

fil = torch.tensor([[1/9,1/9,1/9],[1/9,1/9,1/9],[1/9,1/9,1/9]])
conv_fil = fil.unsqueeze(0).unsqueeze(0)
padded_fil = F.pad(conv_fil, (0, gray_im.shape[0]-fil.shape[0], 0, gray_im.shape[1]-fil.shape[1]))

生成一个填充内核，其中原点位于像素 (1,1)，而不是 (0,0)。它需要在每个方向上移动一个像素。 NumPy 有一个函数 roll 对此很有用，我不知道 Torch 的等价物（我对 Torch 一点也不熟悉）。这应该有效：

fil = torch.tensor([[1/9,1/9,1/9],[1/9,1/9,1/9],[1/9,1/9,1/9]])
padded_fil = fil.unsqueeze(0).unsqueeze(0).numpy()
padded_fil = np.pad(padded_fil, ((0, gray_im.shape[0]-fil.shape[0]), (0, gray_im.shape[1]-fil.shape[1])))
padded_fil = np.roll(padded_fil, -1, axis=(0, 1))
padded_fil = torch.from_numpy(padded_fil)

最后，您的 fftshift 函数应用于空间域图像，导致频域图像（FFT 应用于图像的结果）移动，使得原点位于图像的中间，而不是左上角。这种转变在查看 FFT 的输出时很有用，但在计算卷积时毫无意义。

把这些东西放在一起，卷积现在是：

def complex_multiplication(t1, t2):
  real1, imag1 = t1[:,:,0], t1[:,:,1]
  real2, imag2 = t2[:,:,0], t2[:,:,1]
  return torch.stack([real1 * real2 - imag1 * imag2, real1 * imag2 + imag1 * real2], dim = -1)

fft_im = torch.rfft(gray_im, 2, onesided=False)
fft_fil = torch.rfft(padded_fil, 2, onesided=False)
fft_conv = torch.irfft(complex_multiplication(fft_im, fft_fil), 2, onesided=False)

请注意，您可以进行单侧 FFT 以节省一点计算时间：

fft_im = torch.rfft(gray_im, 2, onesided=True)
fft_fil = torch.rfft(padded_fil, 2, onesided=True)
fft_conv = torch.irfft(complex_multiplication(fft_im, fft_fil), 2, onesided=True, signal_sizes=gray_im.shape)

此处的频域大约是完整 FFT 中的一半大小，但仅省略了冗余部分。卷积结果不变

torch.rfft - 基于 fft 的卷积创建与空间卷积不同的输出

torch.rfft - fft-based convolution creating different output than spatial convolution

python

image-processing

pytorch

fft

convolution