x_fft = torch.fft.rfft(x, dim=1, norm='ortho')

输入输出的形状理解

# x [3584,23,128] [(batch*ts_d), seg_num, d_model]
# Apply FFT along the time dimension
x_fft = torch.fft.rfft(x, dim=1, norm='ortho')
# rfft输出的频率个数=seg_num / 2 + 1
# x_fft [3584,12,128] [(batch*ts_d), seg_num / 2 + 1, d_model]

why

对于实数信号的 FFT（离散傅里叶变换），结果具有共轭对称性，即频谱的后半部分可以由前半部分推导出来。例如，对于一个长度为 ( n ) 的实数序列，其离散傅里叶变换 ( X[k] ) 满足：

\[
X[k] = \overline{X[n-k]}
\]

这意味着 ( k ) 和 ( n-k ) 处的频率信息是冗余的。为了避免重复计算，rfft 只返回非负频率部分，其个数由下式给出：

\[
\text{freq_num} = \left\lfloor \frac{n}{2} \right\rfloor + 1.
\]

具体说明如下：

• 当 \( n \) 为偶数时，比如 \( n=24 \)，唯一的频率个数为 \( \frac{24}{2} + 1 = 13 \)。
• 当 \( n \) 为奇数时，比如 \( n=23 \)，唯一的频率个数为 \( \left\lfloor \frac{23}{2} \right\rfloor + 1 = 11 + 1 = 12 \)。

因此，对于输入张量在指定维度上的长度为 23，rfft 的输出在该维度上的长度就是 12，这样既保留了所有唯一频率分量，也避免了冗余的共轭对称部分。

下面给出一个例子，假设我们有如下数据，其形状为
\[
(\text{batch} \times \text{ts_d},, \text{seg_num},, d_{\text{model}}) = (4,,4,,3)
\]

---

这里我们设定：

• batch = 2
• ts_d = 2
• seg_num = 4（也就是每个信号长度为 4）
• d_model = 3（每个时刻有 3 个特征）

因此，共有 4 个信号（可以认为是 4 个独立样本），每个信号是一个 4×3 的矩阵。比如我们定义数据如下（为了方便说明，这里各个样本的数据取不同值）：

x = tensor([
  [ [ 1,  2,  3],    // 第1个样本，第1个通道的信号： [1, 4, 7, 10] 后面会说明
    [ 4,  5,  6],
    [ 7,  8,  9],
    [10, 11, 12] ],
    
  [ [ 2,  3,  4],
    [ 5,  6,  7],
    [ 8,  9, 10],
    [11, 12, 13] ],
    
  [ [ 3,  4,  5],
    [ 6,  7,  8],
    [ 9, 10, 11],
    [12, 13, 14] ],
    
  [ [ 4,  5,  6],
    [ 7,  8,  9],
    [10, 11, 12],
    [13, 14, 15] ]
])

这里每个最外层的元素对应一个样本（总数为4，即 batch×ts_d），每个样本内部是一个 4×3 的矩阵，其中：

• 第一维（索引0到3）为 seg_num，代表一个信号的时间步数；
• 第二维（索引0到2）为 d_model，对应每个时刻的 3 个特征。

计算过程说明

我们调用

x_fft = torch.fft.rfft(x, dim=1, norm='ortho')

这里说明：

• dim=1 表示对每个样本中的 seg_num 维度进行快速傅里叶变换（FFT），即对每个通道（d_model）独立地对长度为4的序列进行变换。
• norm='ortho' 表示采用正交归一化，即每个变换乘以 ( \frac{1}{\sqrt{N}} )（这里 ( N=4 )）。

1. FFT 后的形状

由于输入沿 dim=1 的长度为 4，为了利用实数输入的共轭对称性质，rfft 只计算非负频率部分，其返回的频率个数为
[
\text{freq_num} = \frac{\text{seg_num}}{2} + 1 = 3.
]
因此，计算后的输出形状为
[
(4,, 3,, 3)
]
即 4 个样本，每个样本有 3 个频率成分，每个频率成分对应 3 个特征（d_model）。

2. 对单个样本、单个通道的计算示例

为了详细说明计算过程，我们以第 0 个样本的第 0 个通道为例。
在该样本中，第 0 个通道的信号为：
[
x[0, :, 0] = [1,, 4,, 7,, 10]
]
我们对这个长度为 4 的实数序列进行 rfft，正交归一化版本的离散傅里叶变换公式为
[
X[k] = \frac{1}{\sqrt{4}} \sum_{n=0}^{3} x[n], e^{-i,2\pi,k,n/4},\quad k=0,1,2.
]

• 频率 ( k=0 ):
  指数部分 ( e^{-i,2\pi,0,n/4} = 1 ) 对所有 ( n ) 都为 1。
  计算：
  [
  X[0] = \frac{1}{2} (1 + 4 + 7 + 10) = \frac{22}{2} = 11.
  ]
  结果为实数 11。

• 频率 ( k=1 ):
  指数部分为：

  • ( n=0 )：( e^{-i,0} = 1 )
  • ( n=1 )：( e^{-i,2\pi/4} = e^{-i,\pi/2} = -i )
  • ( n=2 )：( e^{-i,\pi} = -1 )
  • ( n=3 )：( e^{-i,3\pi/2} = i )

  对应计算：
  [
  \begin{aligned}
  X[1] &= \frac{1}{2} \Bigl[ 1\cdot1 + 4\cdot(-i) + 7\cdot(-1) + 10\cdot i \Bigr] \
  &= \frac{1}{2} \Bigl[ 1 - 4i - 7 + 10i \Bigr] \
  &= \frac{1}{2} \Bigl[ (1-7) + (-4i+10i) \Bigr] \
  &= \frac{1}{2} \Bigl[ -6 + 6i \Bigr] = -3 + 3i.
  \end{aligned}
  ]

• 频率 ( k=2 ):
  此时 ( e^{-i,2\pi,2,n/4} = e^{-i,\pi,n} = (-1)^n )。
  计算：
  [
  \begin{aligned}
  X[2] &= \frac{1}{2} \Bigl[ 1\cdot1 + 4\cdot(-1) + 7\cdot1 + 10\cdot(-1) \Bigr] \
  &= \frac{1}{2} \Bigl[ 1 - 4 + 7 - 10 \Bigr] \
  &= \frac{1}{2} \Bigl[ -6 \Bigr] = -3.
  \end{aligned}
  ]
  注意 ( k=2 )（等于 seg_num/2，当 seg_num 为偶数时）对应的频谱值为实数。

3. 对其他样本和通道

对于形状中每个样本（共 4 个）以及每个通道（3 个），rfft 都是沿着 seg_num 维度分别计算。也就是说，对于第 i 个样本和第 j 个通道，其信号为
[
x[i, :, j] = [x[i,0,j],, x[i,1,j],, x[i,2,j],, x[i,3,j]]
]
计算过程与上面完全类似。如果某个信号的取值不同，则各个频率分量 ( X[k] ) 也会不同；但计算步骤始终按照下式进行：
[
X[k] = \frac{1}{\sqrt{4}} \sum_{n=0}^{3} x[i, n, j], e^{-i,2\pi,k,n/4},\quad k=0,1,2.
]

---

下面我们继续详细说明其他通道的计算过程，并解释如何从输入形状 ([4,4,3]) 得到输出形状 ([4,3,3])（注意：这里输出的第二个维度大小为 ( \frac{4}{2}+1=3 )）。

假设输入张量
[
x = \text{tensor} \Biggl(\begin{array}{ccc}
\begin{bmatrix}
1 & 2 & 3\[4pt]
4 & 5 & 6\[4pt]
7 & 8 & 9\[4pt]
10 & 11 & 12
\end{bmatrix},\quad
\begin{bmatrix}
2 & 3 & 4\[4pt]
5 & 6 & 7\[4pt]
8 & 9 & 10\[4pt]
11 & 12 & 13
\end{bmatrix},\quad
\begin{bmatrix}
3 & 4 & 5\[4pt]
6 & 7 & 8\[4pt]
9 & 10 & 11\[4pt]
12 & 13 & 14
\end{bmatrix},\quad
\begin{bmatrix}
4 & 5 & 6\[4pt]
7 & 8 & 9\[4pt]
10 & 11 & 12\[4pt]
13 & 14 & 15
\end{bmatrix}
\Biggr)
]
其中：

• 最外层维度大小为 4，代表 ( \text{batch} \times \text{ts_d} = 4 ) 个样本；
• 每个样本内部为一个 (4 \times 3) 的矩阵，4 表示 seg_num（时间步数），3 表示 d_model（每个时间步的特征数）。

我们在调用

x_fft = torch.fft.rfft(x, dim=1, norm='ortho')

时，沿着 dim=1（即 seg_num 维度）对每个样本内的每个通道分别做实数 FFT。由于输入长度为 4，因此对于每个信号，rfft 只计算 ( k=0,1,2 ) 三个频率成分（因为 ( \text{freq_num} = \frac{4}{2} + 1 = 3 )），同时归一化因子为 ( \frac{1}{\sqrt{4}} = \frac{1}{2} )。

下面以样本 0 为例，详细说明每个通道的计算。

---

样本 0 各通道的 rfft 计算

通道 0：( x[0, :, 0] = [1,, 4,, 7,, 10] )

归一化因子 ( \alpha = \frac{1}{2} )。

1. ( k=0 )
   [
   X[0] = \alpha \cdot (1+4+7+10) = \frac{22}{2} = 11.
   ]

2. ( k=1 )
   指数项：

   • ( n=0: e^{-i,0} = 1 )
   • ( n=1: e^{-i,\pi/2} = -i )
   • ( n=2: e^{-i,\pi} = -1 )
   • ( n=3: e^{-i,3\pi/2} = i )

   计算：
   [
   \begin{aligned}
   X[1] &= \alpha \Bigl[ 1\cdot1 + 4\cdot(-i) + 7\cdot(-1) + 10\cdot i \Bigr] \
   &= \frac{1}{2}\Bigl[ 1 - 4i - 7 + 10i \Bigr] \
   &= \frac{1}{2}\Bigl[ -6 + 6i \Bigr] = -3 + 3i.
   \end{aligned}
   ]

3. ( k=2 )
   此时指数项 ( e^{-i,\pi n} = (-1)^n )：
   [
   X[2] = \alpha \Bigl[ 1\cdot1 + 4\cdot(-1) + 7\cdot1 + 10\cdot(-1) \Bigr] = \frac{1}{2}\Bigl[ 1-4+7-10 \Bigr] = \frac{-6}{2} = -3.
   ]

因此，样本 0 通道 0 的 rfft 结果为
[
[11,\quad -3+3i,\quad -3].
]

---

通道 1：( x[0, :, 1] = [2,, 5,, 8,, 11] )

同样采用归一化因子 ( \frac{1}{2} )。

1. ( k=0 )
   [
   X[0] = \frac{1}{2}(2+5+8+11) = \frac{26}{2} = 13.
   ]

2. ( k=1 )
   使用相同的指数项：
   [
   \begin{aligned}
   X[1] &= \frac{1}{2}\Bigl[ 2\cdot1 + 5\cdot(-i) + 8\cdot(-1) + 11\cdot i \Bigr] \
   &= \frac{1}{2}\Bigl[ 2 - 5i - 8 + 11i \Bigr] \
   &= \frac{1}{2}\Bigl[ (2-8) + ( -5i+11i) \Bigr] \
   &= \frac{1}{2}\Bigl[ -6 + 6i \Bigr] = -3 + 3i.
   \end{aligned}
   ]

3. ( k=2 )
   [
   X[2] = \frac{1}{2}\Bigl[ 2 - 5 + 8 - 11 \Bigr] = \frac{-6}{2} = -3.
   ]

因此，样本 0 通道 1 的 rfft 结果为
[
[13,\quad -3+3i,\quad -3].
]

---

通道 2：( x[0, :, 2] = [3,, 6,, 9,, 12] )

计算过程同理：

1. ( k=0 )
   [
   X[0] = \frac{1}{2}(3+6+9+12) = \frac{30}{2} = 15.
   ]

2. ( k=1 )
   [
   \begin{aligned}
   X[1] &= \frac{1}{2}\Bigl[ 3\cdot1 + 6\cdot(-i) + 9\cdot(-1) + 12\cdot i \Bigr] \
   &= \frac{1}{2}\Bigl[ 3 - 6i - 9 + 12i \Bigr] \
   &= \frac{1}{2}\Bigl[ (3-9) + ( -6i+12i) \Bigr] \
   &= \frac{1}{2}\Bigl[ -6 + 6i \Bigr] = -3 + 3i.
   \end{aligned}
   ]

3. ( k=2 )
   [
   X[2] = \frac{1}{2}\Bigl[ 3 - 6 + 9 - 12 \Bigr] = \frac{-6}{2} = -3.
   ]

因此，样本 0 通道 2 的 rfft 结果为
[
[15,\quad -3+3i,\quad -3].
]

---

对于其他样本

每个样本都会按照同样的方式独立计算每个通道上的 rfft。例如，对于样本 1，其矩阵为
[
\begin{bmatrix}
2 & 3 & 4\[4pt]
5 & 6 & 7\[4pt]
8 & 9 & 10\[4pt]
11 & 12 & 13
\end{bmatrix}
]

• 样本 1 通道 0： 信号为 ([2,5,8,11])

  • ( k=0 )：( \frac{2+5+8+11}{2} = 13 )
  • ( k=1 )：计算得 ( -3+3i )
  • ( k=2 )：计算得 ( -3 )

• 样本 1 通道 1： 信号为 ([3,6,9,12])

  • ( k=0 )：( \frac{3+6+9+12}{2} = 15 )
  • ( k=1 )：( -3+3i )
  • ( k=2 )：( -3 )

• 样本 1 通道 2： 信号为 ([4,7,10,13])

  • ( k=0 )：( \frac{4+7+10+13}{2} = 17 )
  • ( k=1 )：类似计算得 ( -3+3i )（注意因为各项都是等差递增，所以 k=1 与 k=2 计算模式类似，只是整体和不同）
  • ( k=2 )：( -3 )

其余样本（样本 2 和样本 3）的计算过程同理。

---

输出张量形状的来源

1. 输入形状为 ([4,4,3])：

   • 第一维 4 表示有 4 个样本（即 batch*ts_d）。
   • 第二维 4 表示每个样本有 4 个时间步（seg_num）。
   • 第三维 3 表示每个时间步有 3 个特征（d_model）。

2. FFT 变换沿着 dim=1：
   当对长度为 4 的实数信号进行 rfft 时，只需计算非负频率部分，个数为
   [
   \frac{4}{2} + 1 = 3.
   ]
   因此，每个样本中沿时间步维度转换后得到的频谱维度大小为 3，而 d_model（3）保持不变。

3. 最终输出形状为 ([4,3,3])：

   • 第一维：样本数仍为 4；
   • 第二维：非负频率个数 3（对应 ( k=0,1,2 )）；
   • 第三维：d_model 3，即每个频率分量仍有 3 个特征值。

所以，通过调用

x_fft = torch.fft.rfft(x, dim=1, norm='ortho')

每个样本中，每个通道都得到一个长度为 3 的复数数组，最终拼成形状为 ([4,3,3]) 的输出。

---

总结

• 我们对输入张量中每个样本（形状为 (4 \times 3)）沿时间步（长度为 4）维度进行 rfft，归一化因子为 (1/2)。
• 对于每个通道的计算，举例样本 0 各通道的计算过程分别为：
  • 通道 0 ([1,4,7,10]) 得到 ([11,\ -3+3i,\ -3])；
  • 通道 1 ([2,5,8,11]) 得到 ([13,\ -3+3i,\ -3])；
  • 通道 2 ([3,6,9,12]) 得到 ([15,\ -3+3i,\ -3])。
• 由于 rfft 只保留非负频率，原来 seg_num=4 经过 FFT 后变为 3 个频率分量，因此整个输出张量形状变为 ([4,3,3])。

这样就详细说明了如何从一个形状为 ([4,4,3]) 的输入张量，通过 torch.fft.rfft 沿第二维计算得到输出张量的全过程。