新旧torch中傅里叶变换实现（fft）

news2024/10/5 15:34:34

由泰勒级数我们知道，一个函数可以被分解成无穷个幂函数叠加的形式，于是同样地，一个周期函数也可以被分解成多个周期函数叠加，于是自然而然地，三角函数符合这个需求，由傅里叶级数我们可以将周期函数分解成无穷个三角函数叠加的形式，以下是具体的定义：
在这里插入图片描述

# a, b 均为一个2维tensor

# torch 1.8 之前（旧）
def ccorr(a, b):
    aa = torch.rfft(a, 1)
    bb = torch.rfft(b, 1)
    m = com_mult(conj(aa), bb)
    out = torch.irfft(m, 1, signal_sizes=(a.shape[-1],))
    return out

# torch 1.8 及1.8 之后（新）
def ccorr(a, b):
    aa = torch.fft.rfft2(a, dim=(-2, -1))
    aaa = torch.stack((aa.real, aa.imag), -1)
    bb = torch.fft.rfft2(b, dim=(-2, -1))
    bbb = torch.stack((bb.real, bb.imag), -1)
    m = com_mult(conj(aaa), bbb)
    cor_m = aa * bb.conj()
    # ifft 仅包含实部 irfft包含实部和虚部
    out = torch.fft.irfft2(torch.complex(m[..., 0], m[..., 1]), a.size())
    correlation = torch.fft.irfft2(cor_m, a.size())
    return out