傅里葉卷積如何在PyTorch中實(shí)現(xiàn)

這篇文章將為大家詳細(xì)講解有關(guān)傅里葉卷積如何在PyTorch中實(shí)現(xiàn)，文章內(nèi)容質(zhì)量較高，因此小編分享給大家做個(gè)參考，希望大家閱讀完這篇文章后對(duì)相關(guān)知識(shí)有一定的了解。

卷積

卷積在數(shù)據(jù)分析中無處不在。幾十年來，它們一直被用于信號(hào)和圖像處理。最近，它們成為現(xiàn)代神經(jīng)網(wǎng)絡(luò)的重要組成部分。如果你處理數(shù)據(jù)的話，你可能會(huì)遇到錯(cuò)綜復(fù)雜的問題。

數(shù)學(xué)上，卷積表示為：

盡管離散卷積在計(jì)算應(yīng)用程序中更為常見，但在本文的大部分內(nèi)容中我將使用連續(xù)形式，因?yàn)槭褂眠B續(xù)變量來證明卷積定理(下面討論)要容易得多。之后，我們將回到離散情況，并使用傅立葉變換在 PyTorch 中實(shí)現(xiàn)它。離散卷積可以看作是連續(xù)卷積的近似，其中連續(xù)函數(shù)離散在規(guī)則網(wǎng)格上。因此，我們不會(huì)為這個(gè)離散的案例重新證明卷積定理。

卷積定理

從數(shù)學(xué)上來說，卷積定理可以這樣描述：

其中的連續(xù)傅里葉變換是(達(dá)到正常化常數(shù)) ：

換句話說，位置空間中的卷積等價(jià)于頻率空間中的直乘。這個(gè)想法是相當(dāng)不直觀的，但是對(duì)于連續(xù)的情況來說，證明卷積定理是驚人的容易。要做到這一點(diǎn)，首先要寫出等式的左邊。

現(xiàn)在切換積分的順序，替換變量(x = y + z) ，并分離兩個(gè)被積函數(shù)。

我們?yōu)槭裁匆P(guān)心這一切？

因?yàn)榭焖俑道锶~變換的算法復(fù)雜度低于卷積。直接卷積運(yùn)算具有復(fù)雜度 O(n^2) ，因?yàn)樵?f 中，我們傳遞 g 中的每個(gè)元素，所以可以在 O(nlogn)時(shí)間內(nèi)計(jì)算出快速傅立葉變換。當(dāng)輸入數(shù)組很大時(shí)，它們比卷積要快得多。在這些情況下，我們可以使用卷積定理計(jì)算頻率空間中的卷積，然后執(zhí)行逆傅里葉變換回到位置空間。

當(dāng)輸入較小時(shí)(例如3x3卷積內(nèi)核) ，直接卷積仍然更快。在機(jī)器學(xué)習(xí)應(yīng)用程序中，使用小內(nèi)核更為常見，因此像 PyTorch 和 Tensorflow 這樣的深度學(xué)習(xí)庫只提供直接卷積的實(shí)現(xiàn)。但是在現(xiàn)實(shí)世界中有很多使用大內(nèi)核的用例，其中傅立葉卷積算法更有效。

PyTorch 實(shí)現(xiàn)

現(xiàn)在，我將演示如何在 PyTorch 中實(shí)現(xiàn)傅里葉卷積函數(shù)。它應(yīng)該模仿 torch.nn.functional.convNd 的功能，并利用 fft，而不需要用戶做任何額外的工作。因此，它應(yīng)該接受三個(gè) Tensors (signal、kernel 和可選 bias)和應(yīng)用于輸入的 padding。從概念上講，這個(gè)函數(shù)的內(nèi)部工作原理是：

def fft_conv(
  signal: Tensor, kernel: Tensor, bias: Tensor = None, padding: int = 0,
) -> Tensor:
  # 1. Pad the input signal & kernel tensors
  # 2. Compute FFT for both signal & kernel
  # 3. Multiply the transformed Tensors together
  # 4. Compute inverse FFT
  # 5. Add bias and return

讓我們按照上面顯示的操作順序逐步構(gòu)建 FFT 卷積。對(duì)于這個(gè)例子，我將構(gòu)建一個(gè)一維傅里葉卷積，但是將其擴(kuò)展到二維和三維卷積是很簡單的。

1. 填充輸入數(shù)組

我們需要確保 signal 和 kernel 在填充之后有相同的大小。應(yīng)用初始填充 signal，然后調(diào)整 kernel 的填充以匹配。

# 1. Pad the input signal & kernel tensors
signal = f.pad(signal, [padding, padding])
kernel_padding = [0, signal.size(-1) - kernel.size(-1)]
padded_kernel = f.pad(kernel, kernel_padding)

注意，我只在一邊填充 kernel。我們希望原始內(nèi)核位于填充數(shù)組的左側(cè)，這樣它就可以與 signal 數(shù)組的開始對(duì)齊。

2. 計(jì)算傅立葉變換

這非常簡單，因?yàn)?n 維 fft 已經(jīng)在 PyTorch 中實(shí)現(xiàn)了。我們簡單地使用內(nèi)置函數(shù)，并計(jì)算沿每個(gè)張量的最后一個(gè)維數(shù)的 FFT。

# 2. Perform fourier convolution
signal_fr = rfftn(signal, dim=-1)
kernel_fr = rfftn(padded_kernel, dim=-1)

3. 變換張量相乘

令人驚訝的是，這是我們功能中最復(fù)雜的部分。這有兩個(gè)原因。(1) PyTorch 卷積運(yùn)行于多維張量上，因此我們的 signal 和 kernel 張量實(shí)際上是三維的。從 PyTorch 文檔中的這個(gè)方程式，我們可以看到矩陣乘法是在前兩個(gè)維度上運(yùn)行的(不包括偏差項(xiàng)) ：

我們將需要包括這個(gè)矩陣乘法，以及對(duì)轉(zhuǎn)換后的維度的直接乘法。

PyTorch 實(shí)際上實(shí)現(xiàn)了互相關(guān)/值方法而不是卷積方法。(TensorFlow 和其他深度學(xué)習(xí)庫也是如此。)互相關(guān)與卷積密切相關(guān)，但有一個(gè)重要的標(biāo)志變化：

與卷積相比，這有效地逆轉(zhuǎn)了核的方向(g)。我們不是手動(dòng)翻轉(zhuǎn)內(nèi)核，而是在傅里葉空間中利用內(nèi)核的共軛復(fù)數(shù)來糾正這個(gè)問題。由于我們不需要?jiǎng)?chuàng)建一個(gè)全新的 Tensor，所以這樣做的速度明顯更快，內(nèi)存效率也更高。(本文末尾的附錄中簡要說明了這種方法的工作原理。)

# 3. Multiply the transformed matrices
 
def complex_matmul(a: Tensor, b: Tensor) -> Tensor:
  """Multiplies two complex-valued tensors."""
  # Scalar matrix multiplication of two tensors, over only the first two dimensions.
  # Dimensions 3 and higher will have the same shape after multiplication.
  scalar_matmul = partial(torch.einsum, "ab..., cb... -> ac...") 
 
  # Compute the real and imaginary parts independently, then manually insert them
  # into the output Tensor. This is fairly hacky but necessary for PyTorch 1.7.0,
  # because Autograd is not enabled for complex matrix operations yet. Not exactly
  # idiomatic PyTorch code, but it should work for all future versions (>= 1.7.0).
  real = scalar_matmul(a.real, b.real) - scalar_matmul(a.imag, b.imag)
  imag = scalar_matmul(a.imag, b.real) + scalar_matmul(a.real, b.imag)
  c = torch.zeros(real.shape, dtype=torch.complex64)
  c.real, c.imag = real, imag
  return c 

# Conjugate the kernel for cross-correlation
kernel_fr.imag *= -1
output_fr = complex_matmul(signal_fr, kernel_fr)

PyTorch 1.7改進(jìn)了對(duì)復(fù)數(shù)的支持，但是在 autograd 中還不支持對(duì)復(fù)數(shù)張量的許多操作。現(xiàn)在，我們必須編寫我們自己的復(fù)雜 matmul 方法作為一個(gè)補(bǔ)丁。雖然不是很理想，但是它確實(shí)有效，并且在未來的版本中不會(huì)出現(xiàn)問題。

4. 計(jì)算逆變換

使用 torch.irfftn 可以直接計(jì)算逆變換，然后裁剪出額外的數(shù)組填充。

# 4. Compute inverse FFT, and remove extra padded values
output = irfftn(output_fr, dim=-1)
output = output[:, :, :signal.size(-1) - kernel.size(-1) + 1]

5. 添加偏執(zhí)項(xiàng)并返回

添加偏差項(xiàng)也很容易。請(qǐng)記住，對(duì)于輸出陣列中的每個(gè)通道，偏置項(xiàng)都有一個(gè)元素，并相應(yīng)地調(diào)整其形狀。

# 5. Optionally, add a bias term before returning.
if bias is not None:
  output += bias.view(1, -1, 1)

將上述代碼整合在一起

為了完整起見，讓我們將所有這些代碼片段編譯成一個(gè)內(nèi)聚函數(shù)。

def fft_conv_1d(
  signal: Tensor, kernel: Tensor, bias: Tensor = None, padding: int = 0,
) -> Tensor:
  """
  Args:
    signal: (Tensor) Input tensor to be convolved with the kernel.
    kernel: (Tensor) Convolution kernel.
    bias: (Optional, Tensor) Bias tensor to add to the output.
    padding: (int) Number of zero samples to pad the input on the last dimension.
  Returns:
    (Tensor) Convolved tensor
  """
  # 1. Pad the input signal & kernel tensors
  signal = f.pad(signal, [padding, padding])
  kernel_padding = [0, signal.size(-1) - kernel.size(-1)]
  padded_kernel = f.pad(kernel, kernel_padding)
 
  # 2. Perform fourier convolution
  signal_fr = rfftn(signal, dim=-1)
  kernel_fr = rfftn(padded_kernel, dim=-1)
 
  # 3. Multiply the transformed matrices
  kernel_fr.imag *= -1
  output_fr = complex_matmul(signal_fr, kernel_fr)
 
  # 4. Compute inverse FFT, and remove extra padded values
  output = irfftn(output_fr, dim=-1)
  output = output[:, :, :signal.size(-1) - kernel.size(-1) + 1]
 
  # 5. Optionally, add a bias term before returning.
  if bias is not None:
    output += bias.view(1, -1, 1)
 
 
  return output

直接卷積測(cè)試

最后，我們將使用 torch.nn.functional.conv1d 來確認(rèn)這在數(shù)值上等同于直接一維卷積。我們?yōu)樗休斎霕?gòu)造隨機(jī)張量，并測(cè)量輸出值的相對(duì)差異。

import torch
import torch.nn.functional as f 
 
torch.manual_seed(1234)
kernel = torch.randn(2, 3, 1025)
signal = torch.randn(3, 3, 4096)
bias = torch.randn(2)
 
y0 = f.conv1d(signal, kernel, bias=bias, padding=512)
y1 = fft_conv_1d(signal, kernel, bias=bias, padding=512)
 
abs_error = torch.abs(y0 - y1)
print(f'\nAbs Error Mean: {abs_error.mean():.3E}')
print(f'Abs Error Std Dev: {abs_error.std():.3E}')
 
# Abs Error Mean: 1.272E-05

考慮到我們使用的是32位精度，每個(gè)元素相差大約1e-5ー相當(dāng)精確！讓我們也執(zhí)行一個(gè)快速的基準(zhǔn)來測(cè)量每個(gè)方法的速度：

from timeit import timeit
direct_time = timeit(
  "f.conv1d(signal, kernel, bias=bias, padding=512)", 
  globals=locals(), 
  number=100
) / 100
fourier_time = timeit(
  "fft_conv_1d(signal, kernel, bias=bias, padding=512)", 
  globals=locals(), 
  number=100
) / 100
print(f"Direct time: {direct_time:.3E} s")
print(f"Fourier time: {fourier_time:.3E} s")
 
# Direct time: 1.523E-02 s
# Fourier time: 1.149E-03 s

測(cè)量的基準(zhǔn)將隨著您使用的機(jī)器而發(fā)生顯著的變化。(我正在用一臺(tái)非常舊的 Macbook Pro 進(jìn)行測(cè)試。)對(duì)于1025的內(nèi)核，傅里葉卷積似乎要快10倍以上。

總結(jié)

我希望這已經(jīng)提供了一個(gè)徹底的介紹傅里葉卷積。我認(rèn)為這是一個(gè)非常酷的技巧，在現(xiàn)實(shí)世界中有很多應(yīng)用程序可以使用它。我也喜歡數(shù)學(xué)，所以看到編程和純數(shù)學(xué)的結(jié)合是很有趣的。歡迎和鼓勵(lì)所有的評(píng)論和建設(shè)性的批評(píng)，如果你喜歡這篇文章，請(qǐng)鼓掌！

附錄：

卷積 vs. 互相關(guān)

在本文的前面，我們通過在傅里葉空間中取得內(nèi)核的互相關(guān)共軛復(fù)數(shù)來實(shí)現(xiàn)。這實(shí)際上顛倒了 kernel 的方向，現(xiàn)在我想演示一下為什么會(huì)這樣。首先，記住卷積和互相關(guān)的公式：

然后，讓我們來看看 g(x) 的傅里葉變換：

注意，g(x)是實(shí)值的，所以它不受共軛復(fù)數(shù)變化的影響。然后，更改變量(y =-x)并簡化表達(dá)式。

關(guān)于傅里葉卷積如何在PyTorch中實(shí)現(xiàn)就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，可以學(xué)到更多知識(shí)。如果覺得文章不錯(cuò)，可以把它分享出去讓更多的人看到。